Update app.py

app.py

@@ -18,6 +18,7 @@ from matplotlib.ticker import MaxNLocator
 import gradio as gr
 import tempfile
 import shutil
+import subprocess
 
 # Initialize models and other global variables
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -76,48 +77,79 @@ def alignFace(img):
     new_img = cv2.warpAffine(img_raw, rotation_matrix, (width, height))
     return new_img
 
+def extract_frames(video_path, output_folder, fps):
+    os.makedirs(output_folder, exist_ok=True)
+    command = [
+        'ffmpeg',
+        '-i', video_path,
+        '-vf', f'fps={fps}',
+        f'{output_folder}/frame_%04d.jpg'
+    ]
+    try:
+        subprocess.run(command, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error extracting frames: {e}")
+        print(f"FFmpeg output: {e.output}")
+        raise
+
 def extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired_fps):
-    video = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
-    if not video.isOpened():
-        print(f"Error: Could not open video file at {video_path}")
-        return {}, {}, desired_fps, 0
-    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-    original_fps = video.get(cv2.CAP_PROP_FPS)
-    if frame_count == 0:
-        print(f"Error: Video file at {video_path} appears to be empty")
-        return {}, {}, desired_fps, 0
+    print(f"Processing video: {video_path}")
+    
+    # Extract frames using FFmpeg
+    frames_folder = os.path.join(os.path.dirname(aligned_faces_folder), 'extracted_frames')
+    extract_frames(video_path, frames_folder, desired_fps)
+    
+    # Get video info
+    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_count = int(ffprobe_output[0])
+    original_fps = eval(ffprobe_output[1])
+    
+    print(f"Total frames: {frame_count}, Original FPS: {original_fps}, Desired FPS: {desired_fps}")
+
     embeddings_by_frame = {}
     emotions_by_frame = {}
 
-    for frame_num in range(0, frame_count, int(original_fps / desired_fps)):
-        video.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
-        ret, frame = video.read()
-        if not ret or frame is None or frame.size == 0:
-            print(f"Skipping frame {frame_num}: Frame is empty or couldn't be read")
-            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
+    for frame_file in 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)
+            
+            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
 
-    video.release()
     return embeddings_by_frame, emotions_by_frame, desired_fps, original_fps
 
 def cluster_embeddings(embeddings):
@@ -264,7 +296,7 @@ def plot_emotion(df, emotion):
     ax.xaxis.set_major_locator(MaxNLocator(nbins=100))
     ticks = ax.get_xticks()
     ax.set_xticklabels([df['Timecode'].iloc[int(tick)] if tick >= 0 and tick < len(df) else '' for tick in ticks], rotation=90, ha='right')
-    plt.tight_layout()
+plt.tight_layout()
     return fig
 
 def process_video(video_path, num_anomalies, num_components, desired_fps, batch_size, progress=gr.Progress()):
@@ -275,7 +307,10 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         os.makedirs(organized_faces_folder, exist_ok=True)
 
         progress(0.1, "Extracting and aligning faces")
-        embeddings_by_frame, emotions_by_frame, _, original_fps = extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired_fps)
+        try:
+            embeddings_by_frame, emotions_by_frame, _, original_fps = extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired_fps)
+        except Exception as e:
+            return f"Error extracting faces: {str(e)}", None, None, None, None
 
         if not embeddings_by_frame:
             return "No faces were extracted from the video.", None, None, None, None
@@ -292,11 +327,17 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
 
         progress(0.6, "Performing anomaly detection")
         feature_columns = [col for col in df.columns if col not in ['Frame', 'Timecode', 'Time (Minutes)', 'Embedding_Index']]
-        anomalies_all, anomaly_scores_all, top_indices_all, _ = lstm_anomaly_detection(df[feature_columns].values, feature_columns, num_anomalies=num_anomalies, batch_size=batch_size)
+        try:
+            anomalies_all, anomaly_scores_all, top_indices_all, _ = 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
         
         progress(0.8, "Generating plots")
-        anomaly_plot = plot_anomaly_scores(df, anomaly_scores_all, top_indices_all, "All Features")
-        emotion_plots = [plot_emotion(df, emotion) for emotion in ['fear', 'sad', 'angry']]
+        try:
+            anomaly_plot = plot_anomaly_scores(df, anomaly_scores_all, top_indices_all, "All Features")
+            emotion_plots = [plot_emotion(df, emotion) for emotion in ['fear', 'sad', 'angry']]
+        except Exception as e:
+            return f"Error generating plots: {str(e)}", None, None, None, None
 
         progress(0.9, "Preparing results")
         results = f"Top {num_anomalies} anomalies (All Features):\n"
@@ -327,4 +368,5 @@ iface = gr.Interface(
     description="Upload a video to detect anomalies in facial expressions and emotions. Adjust parameters as needed."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()