go back to 3 step appoach, increase model confidence, use ffmpeg

app.py

@@ -11,7 +11,7 @@ from handlers import video_handler as vh
 model_path = "yolov8n.pt"  # YOLOv8 model path
 
 
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=400)
 def process_video(video_file):
     """
     Processes the uploaded video file by extracting key frames, cropping them, and generating a processed video.
@@ -50,17 +50,21 @@ def process_video(video_file):
     status_message = "Extracting frames. Please wait...!"
     yield status_message, None
 
-    # Step 1: Extract frames with 18 fps to reduce no of processed frames
-    frame_rate = 18
-    vh.extract_frames_by_rate(video_path, all_frames_folder, frame_rate)
+    # Step 1: Extract all frames
+    vh.extract_frames_by_rate(video_path, all_frames_folder, original_fps)
 
-    status_message = "Cropping key frames. Please wait...!"
+    #testing step 2 - extract key frames
+    status_message = "Extracting key frames. Please wait...!"
     yield status_message, None
+    fh.extract_key_frames(all_frames_folder, key_frames_folder, original_fps, model_path)
+    #testing
 
+    status_message = "Cropping key frames. Please wait...!"
+    yield status_message, None
     # Ignore step 2 (extract key frame), do Step 3: Crop key frames based on object detection
     target_resolution = (360, 640)  # Output resolution (9:16)
-    #fh.crop_preserve_key_objects(key_frames_folder, cropped_frames_folder, model_path, target_resolution)
-    fh.crop_preserve_key_objects(all_frames_folder, cropped_frames_folder, model_path, target_resolution)
+    fh.crop_preserve_key_objects(key_frames_folder, cropped_frames_folder, model_path, target_resolution)
+    #fh.crop_preserve_key_objects(all_frames_folder, cropped_frames_folder, model_path, target_resolution)
 
     status_message = "Generating final video. Please wait...!"
     yield status_message, None
@@ -72,7 +76,6 @@ def process_video(video_file):
     status_message = "Processing complete!"
     yield status_message, processed_video_path
 
-
 # Gradio Blocks UI
 with gr.Blocks() as demo:
     gr.Markdown("## Generate short video for your football match")

handlers/frame_handler_yolo.py

@@ -47,7 +47,7 @@ def extract_key_frames(input_folder, key_frames_folder, original_fps, model_path
     # Load YOLO model once
     model = YOLO(model_path)
 
-    # Maintain last 30 non-key frames for reclassification
+    # Maintain last non-key frames for reclassification, max = original_fps
     previous_nonkey_frames = deque(maxlen=original_fps)
     processed_key_frames = set()
     last_frame_was_key = False
@@ -68,16 +68,16 @@ def extract_key_frames(input_folder, key_frames_folder, original_fps, model_path
         if counter % 1000 == 0:
             print(f"Processed {counter} frames.")
         # Run YOLO inference
-        results = model.predict(frame, conf=0.3, verbose=False)
+        results = model.predict(frame, conf=0.7, verbose=False)
 
         # Check if a football (sports ball) is detected
         ball_detected = any(model.names[int(box.cls)] == "sports ball" for box in results[0].boxes)
 
         if ball_detected:
             # TTP: to-do crop the frame
-            # Reclassify up to 30 previous non-key frames
+            # Reclassify up to {original_fps} previous non-key frames
             if not last_frame_was_key:
-                for _ in range(min(len(previous_nonkey_frames), 30)):
+                for _ in range(min(len(previous_nonkey_frames), original_fps)):
                     nonkey_frame_name, nonkey_frame = previous_nonkey_frames.popleft()
                     if nonkey_frame_name not in processed_key_frames:
                         cv2.imwrite(os.path.join(key_frames_folder, nonkey_frame_name), nonkey_frame)
@@ -145,7 +145,7 @@ def crop_preserve_key_objects(input_folder, output_folder, model_path='yolov8n.p
         new_height = int(original_width / target_aspect_ratio)
 
         # YOLO inference
-        results = model.predict(frame, conf=0.5, verbose=False)
+        results = model.predict(frame, conf=0.7, verbose=False)
 
         # Initialize variables
         ball_detected = False

handlers/video_handler.py

@@ -1,4 +1,6 @@
 import os
+import subprocess
+
 import cv2
 
 import functools
@@ -20,61 +22,71 @@ def timer_decorator(func):
 
 @timer_decorator
 def extract_frames_by_rate(video_path, output_folder, frame_rate):
-    """
-    Extracts frames from a video at a specified frame rate.
-
-    Args:
-        video_path (str): Path to the input video file.
-        output_folder (str): Directory to save the extracted frames.
-        frame_rate (int): Number of frames to extract per second of the video.
-    """
-    # Ensure the output directory exists
     if not os.path.exists(output_folder):
         os.makedirs(output_folder)
-
-    # Load the video
-    video = cv2.VideoCapture(video_path)
-
-    # Check if the video is opened successfully
-    if not video.isOpened():
-        print(f"Error: Cannot open video file {video_path}")
-        return
-
-    # Get video properties
-    fps = int(video.get(cv2.CAP_PROP_FPS))  # Frames per second
-    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))  # Total number of frames
-    duration = total_frames / fps  # Duration in seconds
-
-    print(f"Video loaded: {video_path}")
-    print(f"Total Frames: {total_frames}, FPS: {fps}, Duration: {duration:.2f} seconds")
-
-    # Calculate frame interval (in terms of frame number)
-    frame_interval = fps // frame_rate
-
-    # Frame counter
-    frame_count = 0
-    saved_count = 0
-
-    while True:
-        # Read a frame
-        ret, frame = video.read()
-
-        # Break the loop if the video ends
-        if not ret:
-            break
-
-        # Save frame if it matches the frame interval
-        if frame_count % frame_interval == 0:
-            frame_filename = os.path.join(output_folder, f"frame_{saved_count:05d}.jpg")
-            cv2.imwrite(frame_filename, frame)
-            #print(f"Saved: {frame_filename}")
-            saved_count += 1
-
-        frame_count += 1
-
-    # Release video resources
-    video.release()
-    print(f"Extraction complete. Total frames saved: {saved_count}. FPS used to extracted: {frame_rate}")
+    cmd = [
+        'ffmpeg',
+        '-i', video_path,
+        '-vf', f'fps={frame_rate}',
+        os.path.join(output_folder, 'frame_%05d.jpg')
+    ]
+    subprocess.run(cmd, check=True)
+
+# def extract_frames_by_rate(video_path, output_folder, frame_rate):
+#     """
+#     Extracts frames from a video at a specified frame rate.
+#
+#     Args:
+#         video_path (str): Path to the input video file.
+#         output_folder (str): Directory to save the extracted frames.
+#         frame_rate (int): Number of frames to extract per second of the video.
+#     """
+#     # Ensure the output directory exists
+#     if not os.path.exists(output_folder):
+#         os.makedirs(output_folder)
+#
+#     # Load the video
+#     video = cv2.VideoCapture(video_path)
+#
+#     # Check if the video is opened successfully
+#     if not video.isOpened():
+#         print(f"Error: Cannot open video file {video_path}")
+#         return
+#
+#     # Get video properties
+#     fps = int(video.get(cv2.CAP_PROP_FPS))  # Frames per second
+#     total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))  # Total number of frames
+#     duration = total_frames / fps  # Duration in seconds
+#
+#     print(f"Video loaded: {video_path}")
+#     print(f"Total Frames: {total_frames}, FPS: {fps}, Duration: {duration:.2f} seconds")
+#
+#     # Calculate frame interval (in terms of frame number)
+#     frame_interval = fps // frame_rate
+#
+#     # Frame counter
+#     frame_count = 0
+#     saved_count = 0
+#
+#     while True:
+#         # Read a frame
+#         ret, frame = video.read()
+#
+#         # Break the loop if the video ends
+#         if not ret:
+#             break
+#
+#         # Save frame if it matches the frame interval
+#         if frame_count % frame_interval == 0:
+#             frame_filename = os.path.join(output_folder, f"frame_{saved_count:05d}.jpg")
+#             cv2.imwrite(frame_filename, frame)
+#             saved_count += 1
+#
+#         frame_count += 1
+#
+#     # Release video resources
+#     video.release()
+#     print(f"Extraction complete. Total frames saved: {saved_count}. FPS used to extracted: {frame_rate}")
 
 
 @timer_decorator

requirements.txt

@@ -3,3 +3,4 @@ numpy==2.2.3
 opencv_python==4.11.0.86
 spaces==0.32.0
 ultralytics==8.3.64
+ffmpeg-python==0.2.0