Update video_processing.py

video_processing.py

@@ -18,7 +18,7 @@ import torch
 import mediapipe as mp
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.95, 0.95, 0.95], min_face_size=200)
+mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.98, 0.98, 0.98], min_face_size=200, post_process=False)
 
 def extract_frames(video_path, output_folder, desired_fps, progress_callback=None):
     os.makedirs(output_folder, exist_ok=True)
@@ -43,6 +43,21 @@ def extract_frames(video_path, output_folder, desired_fps, progress_callback=Non
     return frame_count, original_fps
 
 
+def is_frontal_face(face, landmarks):
+    if landmarks is None:
+        return False
+    
+    left_eye = landmarks[0]
+    right_eye = landmarks[1]
+    nose = landmarks[2]
+    
+    eye_angle = np.degrees(np.arctan2(right_eye[1] - left_eye[1], right_eye[0] - left_eye[0]))
+    
+    eye_center = ((left_eye[0] + right_eye[0]) / 2, (left_eye[1] + right_eye[1]) / 2)
+    nose_deviation = abs(nose[0] - eye_center[0]) / face.shape[1]
+    
+    return abs(eye_angle) < 10 and nose_deviation < 0.1
+
 def process_frames(frames_folder, faces_folder, frame_count, progress):
     embeddings_by_frame = {}
     posture_scores_by_frame = {}
@@ -60,15 +75,15 @@ def process_frames(frames_folder, faces_folder, frame_count, progress):
             posture_scores_by_frame[frame_num] = posture_score
             posture_landmarks_by_frame[frame_num] = posture_landmarks
 
-            boxes, probs = mtcnn.detect(frame)
+            boxes, probs, landmarks = mtcnn.detect(frame, landmarks=True)
 
             if boxes is not None and len(boxes) > 0 and probs[0] >= 0.99:
                 x1, y1, x2, y2 = [int(b) for b in boxes[0]]
                 face = frame[y1:y2, x1:x2]
-                if face.size > 0:
+
+                if face.size > 0 and is_frontal_face(face, landmarks[0]):
                     face_resized = cv2.resize(face, (160, 160))
                     output_path = os.path.join(faces_folder, f"frame_{frame_num}_face.jpg")
-                    cv2.imwrite(output_path, face_resized)
                     cv2.imwrite(output_path, cv2.cvtColor(face_resized, cv2.COLOR_RGB2BGR))
                     face_paths.append(output_path)
                     embedding = get_face_embedding(face_resized)