Update app.py

app.py

@@ -101,7 +101,7 @@ def extract_frames(video_path, output_folder, fps):
         print(f"FFmpeg stderr: {e.stderr}")
         raise
 
-def extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired_fps):
+def extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired_fps, progress=gr.Progress()):
     print(f"Processing video: {video_path}")
     
     frames_folder = os.path.join(os.path.dirname(aligned_faces_folder), 'extracted_frames')
@@ -123,27 +123,21 @@ def extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired
         if len(ffprobe_output) != 2:
             raise ValueError(f"Unexpected FFprobe output format: {ffprobe_output}")
         
-        frame_count = ffprobe_output[0]
-        frame_rate = ffprobe_output[1]
+        frame_rate, frame_count = ffprobe_output
         
         print(f"Frame count (raw): {frame_count}")
         print(f"Frame rate (raw): {frame_rate}")
         
-        try:
-            frame_count = int(frame_count)
-        except ValueError:
-            print(f"Warning: Could not convert frame count '{frame_count}' to int. Using fallback method.")
-            frame_count = len([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])
-        
         try:
             frac = fractions.Fraction(frame_rate)
             original_fps = float(frac.numerator) / float(frac.denominator)
         except (ValueError, ZeroDivisionError):
             print(f"Warning: Could not convert frame rate '{frame_rate}' to float. Using fallback method.")
-            frame_count = len([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])
             duration_command = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'default=noprint_wrappers=1:nokey=1', video_path]
             duration = float(subprocess.check_output(duration_command, universal_newlines=True).strip())
-            original_fps = frame_count / duration
+            original_fps = int(frame_count) / duration
+        
+        frame_count = int(frame_count)
         
     except subprocess.CalledProcessError as e:
         print(f"Error running FFprobe: {e}")
@@ -157,12 +151,14 @@ def extract_and_align_faces_from_video(video_path, aligned_faces_folder, desired
     embeddings_by_frame = {}
     emotions_by_frame = {}
 
-    for frame_file in sorted(os.listdir(frames_folder)):
+    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) / len(os.listdir(frames_folder)), f"Processing frame {i + 1} of {len(os.listdir(frames_folder))}")
+            
             if frame is None:
                 print(f"Skipping frame {frame_num}: Could not read frame")
                 continue
@@ -368,7 +364,7 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
 
         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)
+            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