Update app.py

app.py

@@ -18,11 +18,14 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
     Priority: YouTube link > Image URL > Uploaded file.
     The sensitivity slider value is passed as the confidence threshold.
     
-    Returns a tuple of 4 items:
-      1. download_file_path (for gr.File)
-      2. image_result (for gr.Image) or None
-      3. video_result (for gr.Video) or None
-      4. status message
+    For video files (mp4, mov, avi, webm), we use streaming mode to obtain annotated frames and encode them into a video.
+    For images, we use the normal prediction and either use the built‑in save_path or plot() method.
+    
+    Returns a tuple:
+      - download_file_path (for gr.File)
+      - image_result (for gr.Image) or None
+      - video_result (for gr.Video) or None
+      - status message
     """
     input_path = None
 
@@ -55,28 +58,68 @@ def process_input(uploaded_file, youtube_link, image_url, sensitivity):
     else:
         return None, None, None, "Please provide an input using one of the methods."
 
-    try:
-        results = model.predict(source=input_path, save=True, conf=sensitivity)
-    except Exception as e:
-        return None, None, None, f"Error running prediction: {e}"
+    # Determine if input is a video (by extension).
+    ext_input = os.path.splitext(input_path)[1].lower()
+    video_exts = [".mp4", ".mov", ".avi", ".webm"]
 
     output_path = None
-    try:
-        if hasattr(results[0], "save_path"):
-            output_path = results[0].save_path
-        else:
-            annotated = results[0].plot()  # returns a numpy array
-            output_path = os.path.join(tempfile.gettempdir(), "annotated.jpg")
-            cv2.imwrite(output_path, annotated)
-    except Exception as e:
-        return None, None, None, f"Error processing the file: {e}"
 
+    if ext_input in video_exts:
+        # Process video using streaming mode.
+        try:
+            # Open video to get properties.
+            cap = cv2.VideoCapture(input_path)
+            if not cap.isOpened():
+                return None, None, None, "Error opening video file."
+            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))
+            cap.release()
+
+            # Use streaming mode to process each frame.
+            frames = []
+            for result in model.predict(source=input_path, stream=True, conf=sensitivity):
+                # result.plot() returns an annotated frame (numpy array)
+                annotated_frame = result.plot()
+                frames.append(annotated_frame)
+            if not frames:
+                return None, None, None, "No detections were returned from video streaming."
+            # Write frames to a temporary video file.
+            temp_video_path = os.path.join(tempfile.gettempdir(), "annotated_video.mp4")
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
+            for frame in frames:
+                out.write(frame)
+            out.release()
+            output_path = temp_video_path
+        except Exception as e:
+            return None, None, None, f"Error processing video: {e}"
+    else:
+        # Process as an image.
+        try:
+            results = model.predict(source=input_path, save=True, conf=sensitivity)
+        except Exception as e:
+            return None, None, None, f"Error running prediction: {e}"
+
+        try:
+            if not results or len(results) == 0:
+                return None, None, None, "No detections were returned."
+            if hasattr(results[0], "save_path"):
+                output_path = results[0].save_path
+            else:
+                annotated = results[0].plot()  # returns a numpy array
+                output_path = os.path.join(tempfile.gettempdir(), "annotated.jpg")
+                cv2.imwrite(output_path, annotated)
+        except Exception as e:
+            return None, None, None, f"Error processing the file: {e}"
+
+    # Clean up temporary input if downloaded.
     if ((youtube_link and youtube_link.strip()) or (image_url and image_url.strip())) and input_path and os.path.exists(input_path):
         os.remove(input_path)
 
-    ext = os.path.splitext(output_path)[1].lower()
-    video_exts = [".mp4", ".mov", ".avi", ".webm"]
-    if ext in video_exts:
+    # Set outputs based on output file extension.
+    ext_output = os.path.splitext(output_path)[1].lower()
+    if ext_output in video_exts:
         image_result = None
         video_result = output_path
     else: