Update app.py

app.py

@@ -6,25 +6,18 @@ import cv2
 import requests
 from ultralytics import YOLO
 
-# Remove extra CLI arguments that Spaces might pass
 sys.argv = [arg for arg in sys.argv if arg != "--import"]
 
-# Load the YOLO11-pose model
 model = YOLO("yolo11n-pose.pt")
 
 def process_input(uploaded_file, youtube_link, image_url, sensitivity):
-    """
-    Process input from Upload, YouTube, or Image URL.
-    Priority: YouTube link > Image URL > Uploaded file.
-    Sensitivity is the confidence threshold.
-    """
     input_path = None
     temp_files = []
 
-    # Priority 1: YouTube link
+    # Input priority handling
     if youtube_link and youtube_link.strip():
         try:
-            from pytubefix import YouTube  # Use pytubefix instead of pytube
+            from pytubefix import YouTube
             yt = YouTube(youtube_link)
             stream = yt.streams.filter(file_extension='mp4', progressive=True).order_by("resolution").desc().first()
             if not stream:
@@ -35,8 +28,6 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
             temp_files.append(input_path)
         except Exception as e:
             return None, None, None, f"Error downloading YouTube video: {str(e)}"
-
-    # Priority 2: Image URL
     elif image_url and image_url.strip():
         try:
             response = requests.get(image_url, stream=True, timeout=10)
@@ -48,14 +39,12 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
             temp_files.append(input_path)
         except Exception as e:
             return None, None, None, f"Error downloading image: {str(e)}"
-
-    # Priority 3: Uploaded file
     elif uploaded_file is not None:
         input_path = uploaded_file.name
     else:
         return None, None, None, "Please provide an input."
 
-    # Process the file
+    # Process file
     ext = os.path.splitext(input_path)[1].lower()
     video_exts = [".mp4", ".mov", ".avi", ".webm"]
     output_path = None
@@ -65,38 +54,50 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
             # Video processing
             cap = cv2.VideoCapture(input_path)
             if not cap.isOpened():
-                return None, None, None, "Error opening video file."
+                return None, None, None, f"Cannot open video file: {input_path}"
             
             fps = cap.get(cv2.CAP_PROP_FPS)
             width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
             height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            
+            if fps <= 0 or width <= 0 or height <= 0:
+                return None, None, None, "Invalid video properties detected."
             
-            # Create output video
             output_path = os.path.join(tempfile.gettempdir(), f"out_{os.urandom(8).hex()}.mp4")
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            # Use H.264 codec which is more widely supported
+            fourcc = cv2.VideoWriter_fourcc(*'avc1')  # Changed from 'mp4v' to 'avc1'
             out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
             
+            if not out.isOpened():
+                return None, None, None, "Failed to initialize video writer."
+            
+            processed_frames = 0
             while True:
                 ret, frame = cap.read()
                 if not ret:
                     break
                     
-                # Convert BGR to RGB for YOLO
+                # Process frame
                 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 results = model.predict(source=frame_rgb, conf=sensitivity)[0]
                 annotated_frame = results.plot()
-                # Convert back to BGR for video writing
                 annotated_frame_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
+                
                 out.write(annotated_frame_bgr)
+                processed_frames += 1
             
             cap.release()
             out.release()
             temp_files.append(output_path)
             
-            if os.path.getsize(output_path) == 0:
-                return None, None, None, "Error: Output video is empty."
+            if processed_frames == 0:
+                return None, None, None, "No frames processed from video."
+                
+            if not os.path.exists(output_path) or os.path.getsize(output_path) < 1024:  # Check if file is too small
+                return None, None, None, f"Output video created but too small ({os.path.getsize(output_path)} bytes) - processing failed."
                 
-            return output_path, None, output_path, "Video processed successfully!"
+            return output_path, None, output_path, f"Video processed successfully! ({processed_frames}/{frame_count} frames)"
 
         else:
             # Image processing
@@ -109,17 +110,16 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
 
     except Exception as e:
         return None, None, None, f"Processing error: {str(e)}"
-
+    
     finally:
-        # Clean up temporary input files (but keep output for download)
-        for f in temp_files[:-1]:  # Exclude output_path
+        for f in temp_files[:-1]:
             if f and os.path.exists(f):
                 try:
                     os.remove(f)
                 except:
                     pass
 
-# Gradio interface remains mostly the same
+# Gradio interface (unchanged)
 with gr.Blocks(css="""
 .result_img > img {
   width: 100%;