Update video_processing.py

video_processing.py

@@ -24,6 +24,74 @@ mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_fac
 mp_face_mesh = mp.solutions.face_mesh
 face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.8)
 
+
+def extract_frames(video_path, output_folder, desired_fps, progress_callback=None):
+    os.makedirs(output_folder, exist_ok=True)
+    clip = VideoFileClip(video_path)
+    original_fps = clip.fps
+    duration = clip.duration
+    total_frames = int(duration * original_fps)
+    step = max(1, original_fps / desired_fps)
+    total_frames_to_extract = int(total_frames / step)
+
+    frame_count = 0
+    for t in np.arange(0, duration, step / original_fps):
+        frame = clip.get_frame(t)
+        cv2.imwrite(os.path.join(output_folder, f"frame_{frame_count:04d}.jpg"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+        frame_count += 1
+        if progress_callback:
+            progress = min(100, (frame_count / total_frames_to_extract) * 100)
+            progress_callback(progress, f"Extracting frame")
+        if frame_count >= total_frames_to_extract:
+            break
+    clip.close()
+    return frame_count, original_fps
+
+
+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')])
+
+    for i, frame_file in enumerate(frame_files):
+        frame_num = int(frame_file.split('_')[1].split('.')[0])
+        frame_path = os.path.join(frames_folder, frame_file)
+        frame = cv2.imread(frame_path)
+
+        if frame is not None:
+            posture_score, posture_landmarks = calculate_posture_score(frame)
+            posture_scores_by_frame[frame_num] = posture_score
+            posture_landmarks_by_frame[frame_num] = posture_landmarks
+
+            boxes, probs = mtcnn.detect(frame)
+
+            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:
+                    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
+
+        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
+
+
 def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
     start_time = time.time()
     output_folder = "output"