Update utils.py

utils.py

@@ -42,6 +42,7 @@ def create_annotated_video(video_path, df, mse_embeddings, largest_cluster, outp
     import torch
     from facenet_pytorch import MTCNN
     import mediapipe as mp
+    import numpy as np
     
     video = cv2.VideoCapture(video_path)
     fps = video.get(cv2.CAP_PROP_FPS)
@@ -49,23 +50,28 @@ def create_annotated_video(video_path, df, mse_embeddings, largest_cluster, outp
     height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
+
+    mp_face_mesh = mp.solutions.face_mesh
+    mp_drawing = mp.solutions.drawing_utils
+    mp_drawing_styles = mp.solutions.drawing_styles
+    face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.5)
+
     frame_number = 0
     while True:
         ret, frame = video.read()
         if not ret:
             break
+
         # Detect face and draw bounding box
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
         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
-        mp_face_mesh = mp.solutions.face_mesh
-        mp_drawing = mp.solutions.drawing_utils
-        mp_drawing_styles = mp.solutions.drawing_styles
-        face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.5)
         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:
@@ -76,11 +82,19 @@ def create_annotated_video(video_path, df, mse_embeddings, largest_cluster, outp
                     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]
+            frame_index = np.where(df['Frame'].values == frame_number)[0][0]
+            if mse_embeddings.ndim == 1:
+                mse = mse_embeddings[frame_index]
+            else:
+                mse = mse_embeddings[frame_index, 0]  # Assuming MSE is in the first column if 2D
             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()
+    out.release()
+    face_mesh.close()