Update utils.py

utils.py

@@ -26,6 +26,7 @@ def add_timecode_to_image(image, timecode):
     font = ImageFont.truetype("arial.ttf", 15)
     draw.text((10, 10), timecode, (255, 0, 0), font=font)
     return np.array(img_pil)
+    
 
 def add_timecode_to_image_body(image, timecode):
     from PIL import Image, ImageDraw, ImageFont
@@ -37,6 +38,39 @@ def add_timecode_to_image_body(image, timecode):
     draw.text((10, 10), timecode, (255, 0, 0), font=font)
     return np.array(img_pil)
 
-def timecode_to_seconds(timecode):
-    h, m, s = map(float, timecode.split(':'))
-    return h * 3600 + m * 60 + s
+def create_annotated_video(video_path, df, mse_embeddings, largest_cluster, output_path):
+    video = cv2.VideoCapture(video_path)
+    fps = video.get(cv2.CAP_PROP_FPS)
+    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    frame_number = 0
+    while True:
+        ret, frame = video.read()
+        if not ret:
+            break
+        # Detect face and draw bounding box
+        boxes, _ = mtcnn.detect(frame)
+        if boxes is not None and len(boxes) > 0:
+            box = boxes[0]
+            cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
+        # Draw facial landmarks
+        face_mesh_results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        if face_mesh_results.multi_face_landmarks:
+            for face_landmarks in face_mesh_results.multi_face_landmarks:
+                mp_drawing.draw_landmarks(
+                    image=frame,
+                    landmark_list=face_landmarks,
+                    connections=mp_face_mesh.FACEMESH_TESSELATION,
+                    landmark_drawing_spec=None,
+                    connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_tesselation_style()
+                )
+        # Add MSE annotation
+        if frame_number in df['Frame'].values:
+            mse = mse_embeddings[df['Frame'] == frame_number].iloc[0]
+            cv2.putText(frame, f"MSE: {mse:.4f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+        out.write(frame)
+        frame_number += 1
+    video.release()
+    out.release()