Update video_processing.py

video_processing.py

@@ -47,7 +47,6 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
     embeddings_by_frame = {}
     posture_scores_by_frame = {}
     posture_landmarks_by_frame = {}
-    facial_landmarks_by_frame = {}
     aligned_face_paths = []
     frame_files = sorted([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])
 
@@ -68,23 +67,16 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
                 face = frame[y1:y2, x1:x2]
                 if face.size > 0:
                     face_rgb = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
-                    results = face_mesh.process(face_rgb)
-                    if results.multi_face_landmarks:
-                        facial_landmarks_by_frame[frame_num] = results.multi_face_landmarks[0]
-                        if is_frontal_face(results.multi_face_landmarks[0].landmark):
-                            aligned_face = face
-
-                            if aligned_face is not None:
-                                aligned_face_resized = cv2.resize(aligned_face, (160, 160))
-                                output_path = os.path.join(aligned_faces_folder, f"frame_{frame_num}_face.jpg")
-                                cv2.imwrite(output_path, aligned_face_resized)
-                                aligned_face_paths.append(output_path)
-                                embedding = get_face_embedding(aligned_face_resized)
-                                embeddings_by_frame[frame_num] = embedding
+                    face_resized = cv2.resize(face_rgb, (160, 160))
+                    output_path = os.path.join(faces_folder, f"frame_{frame_num}_face.jpg")
+                    cv2.imwrite(output_path, face_resized_resized)
+                    face_paths.append(output_path)
+                    embedding = get_face_embedding(face_resized)
+                    embeddings_by_frame[frame_num] = embedding
 
         progress((i + 1) / len(frame_files), f"Processing frame {i + 1} of {len(frame_files)}")
 
-    return embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, aligned_face_paths, facial_landmarks_by_frame
+    return embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, aligned_face_paths
 
 
 def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
@@ -93,9 +85,9 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
     os.makedirs(output_folder, exist_ok=True)
 
     with tempfile.TemporaryDirectory() as temp_dir:
-        aligned_faces_folder = os.path.join(temp_dir, 'aligned_faces')
+        faces_folder = os.path.join(temp_dir, 'faces')
         organized_faces_folder = os.path.join(temp_dir, 'organized_faces')
-        os.makedirs(aligned_faces_folder, exist_ok=True)
+        os.makedirs(faces_folder, exist_ok=True)
         os.makedirs(organized_faces_folder, exist_ok=True)
 
         clip = VideoFileClip(video_path)
@@ -112,12 +104,12 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
 
         progress(1, "Frame extraction complete")
         progress(0.3, "Processing frames")
-        embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, aligned_face_paths, facial_landmarks_by_frame = process_frames(
-            frames_folder, aligned_faces_folder,
+        embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, face_paths = process_frames(
+            frames_folder, faces_folder,
             frame_count,
             progress)
 
-        if not aligned_face_paths:
+        if not face_paths:
             raise ValueError("No faces were extracted from the video.")
 
         progress(0.6, "Clustering faces")
@@ -129,10 +121,10 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
         cluster_by_frame = {frame_num: cluster for frame_num, cluster in zip(embeddings_by_frame.keys(), clusters)}
 
         progress(0.65, "Organizing faces")
-        organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder)
+        organize_faces_by_person(embeddings_by_frame, clusters, faces_folder, organized_faces_folder)
 
         progress(0.7, "Saving person data")
-        df, largest_cluster = save_person_data_to_csv(embeddings_by_frame, clusters, desired_fps,
+        df, largest_cluster = save_person_data(embeddings_by_frame, clusters, desired_fps,
                                                       original_fps, temp_dir, video_duration)
 
         df['Seconds'] = df['Timecode'].apply(
@@ -223,7 +215,7 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
 
         anomaly_faces_embeddings = []
         for frame in anomaly_frames_embeddings:
-            face_path = os.path.join(aligned_faces_folder, f"frame_{frame}_face.jpg")
+            face_path = os.path.join(faces_folder, f"frame_{frame}_face.jpg")
             if os.path.exists(face_path):
                 face_img = cv2.imread(face_path)
                 if face_img is not None:
@@ -265,28 +257,14 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
             face_samples["most_frequent"],
             anomaly_faces_embeddings,
             anomaly_frames_posture_images,
-            aligned_faces_folder,
+            faces_folder,
             frames_folder,
             stacked_heatmap
               
         )
 
 
-def is_frontal_face(landmarks, threshold=60):
-    nose_tip = landmarks[4]
-    left_chin = landmarks[234]
-    right_chin = landmarks[454]
-    nose_to_left = [left_chin.x - nose_tip.x, left_chin.y - nose_tip.y]
-    nose_to_right = [right_chin.x - nose_tip.x, right_chin.y - nose_tip.y]
-    dot_product = nose_to_left[0] * nose_to_right[0] + nose_to_left[1] * nose_to_right[1]
-    magnitude_left = math.sqrt(nose_to_left[0] ** 2 + nose_to_left[1] ** 2)
-    magnitude_right = math.sqrt(nose_to_right[0] ** 2 + nose_to_right[1] ** 2)
-    cos_angle = dot_product / (magnitude_left * magnitude_right)
-    angle = math.acos(cos_angle)
-    angle_degrees = math.degrees(angle)
-    return abs(180 - angle_degrees) < threshold
-
-def save_person_data_to_csv(embeddings_by_frame, clusters, desired_fps, original_fps, output_folder, video_duration):
+def save_person_data(embeddings_by_frame, clusters, desired_fps, original_fps, output_folder, video_duration):
     person_data = {}
 
     for (frame_num, embedding), cluster in zip(embeddings_by_frame.items(), clusters):
@@ -319,7 +297,7 @@ def save_person_data_to_csv(embeddings_by_frame, clusters, desired_fps, original
 
     return df, largest_cluster
 
-def get_all_face_samples(organized_faces_folder, output_folder, largest_cluster, max_samples=100):
+def get_all_face_samples(organized_faces_folder, output_folder, largest_cluster, max_samples=200):
     face_samples = {"most_frequent": [], "others": []}
     for cluster_folder in sorted(os.listdir(organized_faces_folder)):
         if cluster_folder.startswith("person_"):