Update app.py

app.py

@@ -362,33 +362,37 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         os.makedirs(aligned_faces_folder, exist_ok=True)
         os.makedirs(organized_faces_folder, exist_ok=True)
 
-        progress(0.1, "Extracting and aligning faces")
-        try:
-            embeddings_by_frame, emotions_by_frame, _, original_fps = extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired_fps, progress)
-        except Exception as e:
-            return f"Error extracting faces: {str(e)}", None, None, None, None, None, None
+        progress(0.1, "Extracting frames")
+        frames_folder = os.path.join(temp_dir, 'extracted_frames')
+        extract_frames(video_path, frames_folder, desired_fps)
+
+        progress(0.2, "Getting video info")
+        frame_count, original_fps = get_video_info(video_path)
+
+        progress(0.3, "Processing frames")
+        embeddings_by_frame, emotions_by_frame = process_frames(frames_folder, aligned_faces_folder, frame_count, progress)
 
         if not embeddings_by_frame:
             return "No faces were extracted from the video.", None, None, None, None, None, None
 
-        progress(0.3, "Clustering embeddings")
+        progress(0.6, "Clustering embeddings")
         embeddings = list(embeddings_by_frame.values())
         clusters = cluster_embeddings(embeddings)
 
-        progress(0.4, "Organizing faces")
+        progress(0.7, "Organizing faces")
         organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder)
 
-        progress(0.5, "Saving person data")
+        progress(0.8, "Saving person data")
         df, largest_cluster = save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, desired_fps, original_fps, temp_dir, num_components)
 
-        progress(0.6, "Performing anomaly detection")
+        progress(0.9, "Performing anomaly detection")
         feature_columns = [col for col in df.columns if col not in ['Frame', 'Timecode', 'Time (Minutes)', 'Embedding_Index']]
         try:
             anomalies_all, anomaly_scores_all, top_indices_all, anomalies_comp, anomaly_scores_comp, top_indices_comp, _ = lstm_anomaly_detection(df[feature_columns].values, feature_columns, num_anomalies=num_anomalies, batch_size=batch_size)
         except Exception as e:
             return f"Error in anomaly detection: {str(e)}", None, None, None, None, None, None
         
-        progress(0.8, "Generating plots")
+        progress(0.95, "Generating plots")
         try:
             anomaly_plot_all = plot_anomaly_scores(df, anomaly_scores_all, top_indices_all, "All Features")
             anomaly_plot_comp = plot_anomaly_scores(df, anomaly_scores_comp, top_indices_comp, "Components Only")
@@ -397,7 +401,7 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         except Exception as e:
             return f"Error generating plots: {str(e)}", None, None, None, None, None, None
 
-        progress(0.9, "Preparing results")
+        progress(1.0, "Preparing results")
         results = f"Top {num_anomalies} anomalies (All Features):\n"
         results += "\n".join([f"{score:.4f} at {timecode}" for score, timecode in 
                               zip(anomaly_scores_all[top_indices_all], df['Timecode'].iloc[top_indices_all].values)])
@@ -405,15 +409,72 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         results += "\n".join([f"{score:.4f} at {timecode}" for score, timecode in 
                               zip(anomaly_scores_comp[top_indices_comp], df['Timecode'].iloc[top_indices_comp].values)])
 
-        # Add top emotion scores to results
         for emotion in ['fear', 'sad', 'angry']:
             top_indices = np.argsort(df[emotion].values)[-num_anomalies:][::-1]
             results += f"\n\nTop {num_anomalies} {emotion.capitalize()} Scores:\n"
             results += "\n".join([f"{df[emotion].iloc[i]:.4f} at {df['Timecode'].iloc[i]}" for i in top_indices])
 
-        progress(1.0, "Complete")
         return results, anomaly_plot_all, anomaly_plot_comp, components_plot, *emotion_plots
 
+def get_video_info(video_path):
+    ffprobe_command = [
+        'ffprobe',
+        '-v', 'error',
+        '-select_streams', 'v:0',
+        '-count_packets',
+        '-show_entries', 'stream=nb_read_packets,r_frame_rate',
+        '-of', 'csv=p=0',
+        video_path
+    ]
+    ffprobe_output = subprocess.check_output(ffprobe_command, universal_newlines=True).strip().split(',')
+    frame_rate, frame_count = ffprobe_output
+    
+    frac = fractions.Fraction(frame_rate)
+    original_fps = float(frac.numerator) / float(frac.denominator)
+    frame_count = int(frame_count)
+    
+    return frame_count, original_fps
+
+def process_frames(frames_folder, aligned_faces_folder, frame_count, progress):
+    embeddings_by_frame = {}
+    emotions_by_frame = {}
+
+    for i, frame_file in enumerate(sorted(os.listdir(frames_folder))):
+        if frame_file.endswith('.jpg'):
+            frame_num = int(frame_file.split('_')[1].split('.')[0])
+            frame_path = os.path.join(frames_folder, frame_file)
+            frame = cv2.imread(frame_path)
+            
+            progress((i + 1) / frame_count, f"Processing frame {i + 1} of {frame_count}")
+            
+            if frame is None:
+                print(f"Skipping frame {frame_num}: Could not read frame")
+                continue
+
+            try:
+                boxes, probs = mtcnn.detect(frame)
+                if boxes is not None and len(boxes) > 0:
+                    box = boxes[0]
+                    if probs[0] >= 0.99:
+                        x1, y1, x2, y2 = [int(b) for b in box]
+                        face = frame[y1:y2, x1:x2]
+                        if face.size == 0:
+                            print(f"Skipping frame {frame_num}: Detected face region is empty")
+                            continue
+                        aligned_face = alignFace(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)
+                            embedding, emotion = get_face_embedding_and_emotion(aligned_face_resized)
+                            embeddings_by_frame[frame_num] = embedding
+                            emotions_by_frame[frame_num] = emotion
+            except Exception as e:
+                print(f"Error processing frame {frame_num}: {str(e)}")
+                continue
+
+    return embeddings_by_frame, emotions_by_frame
+
 # Gradio interface
 iface = gr.Interface(
     fn=process_video,