Update app.py

app.py

@@ -2,12 +2,12 @@ import cv2
 import torch
 import gradio as gr
 import numpy as np
-from ultralytics import YOLO, __version__ as ultralytics_version  # Import ultralytics version
+from ultralytics import YOLO, __version__ as ultralytics_version
 
 # Debug: Check environment
 print(f"Torch version: {torch.__version__}")
 print(f"Gradio version: {gr.__version__}")
-print(f"Ultralytics version: {ultralytics_version}")  # Corrected version access
+print(f"Ultralytics version: {ultralytics_version}")
 print(f"CUDA available: {torch.cuda.is_available()}")
 
 # Load YOLOv8 model
@@ -19,39 +19,50 @@ def process_video(video):
     if video is None:
         return "Error: No video uploaded"
     
+    # Open the input video
     cap = cv2.VideoCapture(video)
     if not cap.isOpened():
         return "Error: Could not open video file"
     
-    frame_width, frame_height = 320, 240  # Smaller resolution
+    # Get input video properties
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
+    # Use original resolution to avoid resizing issues (optional: keep 320x240 if needed)
+    # frame_width, frame_height = 320, 240
+    print(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames")
+    
+    # Set up video writer
     output_path = "processed_output.mp4"
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
     
-    frame_count = 1
-    frame_skip = 5  # Process every 5th frame
-    max_frames = 100  # Limit for testing
+    frame_count = 0
     
     while True:
         ret, frame = cap.read()
-        if not ret or frame_count > max_frames:
+        if not ret:
             break
             
         frame_count += 1
-        if frame_count % frame_skip != 0:
-            continue
-            
-        frame = cv2.resize(frame, (frame_width, frame_height))
-        print(f"Processing frame {frame_count}")
+        print(f"Processing frame {frame_count}/{total_frames}")
         
+        # Optional: Resize if needed (remove if using original resolution)
+        # frame = cv2.resize(frame, (frame_width, frame_height))
+        
+        # Run YOLOv8 inference
         results = model(frame)
         annotated_frame = results[0].plot()
+        
+        # Write the annotated frame to the output video
         out.write(annotated_frame)
     
+    # Release resources
     cap.release()
     out.release()
+    print(f"Output video saved as {output_path}")
     return output_path
 
 # Gradio interface