Update app.py

app.py

@@ -5,7 +5,8 @@ from huggingface_hub import snapshot_download
 import os
 import cv2
 import numpy as np
-from tqdm import tqdm  # For progress bar
+from tqdm import tqdm
+import tempfile
 
 # Function to load the model
 def load_model(repo_id):
@@ -15,7 +16,7 @@ def load_model(repo_id):
     detection_model = YOLO(path, task="detect")
     return detection_model
 
-# Function to predict an image
+# Function to process an image
 def predict_image(pilimg, conf_threshold, iou_threshold):
     """Process an image with user-defined thresholds."""
     try:
@@ -26,9 +27,9 @@ def predict_image(pilimg, conf_threshold, iou_threshold):
     except Exception as e:
         return f"Error processing image: {e}"
 
-# Function to predict a video with progress tracking
+# Function to process a video
 def predict_video(video_file, conf_threshold, iou_threshold, start_time, end_time):
-    """Process a video with user-defined thresholds and time range."""
+    """Process a video and return the path for displaying."""
     cap = cv2.VideoCapture(video_file)
 
     if not cap.isOpened():
@@ -39,9 +40,11 @@ def predict_video(video_file, conf_threshold, iou_threshold, start_time, end_tim
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    output_path = "output_video.mp4"
+    # Use a temporary file to store the processed video
+    temp_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    output_path = temp_video_file.name
 
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     start_frame = int(start_time * fps) if start_time else 0
     end_frame = int(end_time * fps) if end_time else total_frames
 
@@ -78,6 +81,7 @@ detection_model = load_model(REPO_ID)
 with gr.Blocks() as demo:
     gr.Markdown("## Pangolin and Axolotl Detection")
 
+    # Image Processing Tab
     with gr.Tab("Image Input"):
         img_input = gr.Image(type="pil", label="Upload an Image")
         conf_slider_img = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold")
@@ -91,13 +95,14 @@ with gr.Blocks() as demo:
             outputs=img_output
         )
 
+    # Video Processing Tab
     with gr.Tab("Video Input"):
         video_input = gr.Video(label="Upload a Video")
         conf_slider_video = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold")
         iou_slider_video = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="IoU Threshold")
         start_time = gr.Number(value=0, label="Start Time (seconds)")
         end_time = gr.Number(value=0, label="End Time (seconds, 0 for full video)")
-        video_output = gr.File(label="Download Processed Video")
+        video_output = gr.Video(label="Processed Video")
         video_submit = gr.Button("Process Video")
 
         video_submit.click(