Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import os
 import matplotlib.pyplot as plt
 from ultralytics import YOLO, __version__ as ultralytics_version
+import uuid
 
 # Debug: Check environment
 print(f"Torch version: {torch.__version__}")
@@ -19,7 +20,8 @@ model = YOLO('./data/best.pt').to(device)
 
 def process_video(video, output_folder="detected_frames", plot_graphs=False):
     if video is None:
-        return "Error: No video uploaded"
+        yield "Error: No video uploaded", []
+        return
     
     # Create output folder if it doesn't exist
     if not os.path.exists(output_folder):
@@ -27,13 +29,15 @@ def process_video(video, output_folder="detected_frames", plot_graphs=False):
     
     cap = cv2.VideoCapture(video)
     if not cap.isOpened():
-        return "Error: Could not open video file"
+        yield "Error: Could not open video file", []
+        return
     
     frame_width, frame_height = 320, 240  # Smaller resolution
     frame_count = 0
     frame_skip = 5  # Process every 5th frame
     max_frames = 100  # Limit for testing
     confidence_scores = []  # Store confidence scores for plotting
+    detected_frame_paths = []  # Store paths of frames with detections
     
     while True:
         ret, frame = cap.read()
@@ -49,21 +53,24 @@ def process_video(video, output_folder="detected_frames", plot_graphs=False):
         
         # Run YOLOv8 inference
         results = model(frame)
-        annotated_frame = results[0].plot()
         
-        # Save annotated frame
-        frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.jpg")
-        cv2.imwrite(frame_filename, annotated_frame)
-        
-        # Collect confidence scores for plotting
-        if results[0].boxes is not None:
+        # Save and yield frame if objects are detected
+        if results[0].boxes is not None and len(results[0].boxes) > 0:
+            annotated_frame = results[0].plot()
+            frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.jpg")
+            cv2.imwrite(frame_filename, annotated_frame)
+            detected_frame_paths.append(frame_filename)
+            
+            # Collect confidence scores for plotting
             confs = results[0].boxes.conf.cpu().numpy()
             confidence_scores.extend(confs)
+            
+            # Yield current status and gallery
+            yield f"Processed frame {frame_count} with detections", detected_frame_paths[:]
     
     cap.release()
     
     # Generate confidence score plot if requested
-    graph_path = None
     if plot_graphs and confidence_scores:
         plt.figure(figsize=(10, 5))
         plt.hist(confidence_scores, bins=20, color='blue', alpha=0.7)
@@ -73,21 +80,33 @@ def process_video(video, output_folder="detected_frames", plot_graphs=False):
         graph_path = os.path.join(output_folder, "confidence_histogram.png")
         plt.savefig(graph_path)
         plt.close()
+        detected_frame_paths.append(graph_path)
     
-    return f"Frames saved in {output_folder}. {f'Graph saved as {graph_path}' if graph_path else ''}"
+    # Final yield with all results
+    status = f"Saved {len(detected_frame_paths)} frames with detections in {output_folder}. {f'Graph saved as {graph_path}' if plot_graphs and confidence_scores else ''}"
+    yield status, detected_frame_paths
 
 # Gradio interface
-iface = gr.Interface(
-    fn=process_video,
-    inputs=[
-        gr.Video(label="Upload Video"),
-        gr.Textbox(label="Output Folder", value="detected_frames"),
-        gr.Checkbox(label="Generate Confidence Score Graph", value=False)
-    ],
-    outputs=gr.Text(label="Status"),
-    title="YOLOv8 Object Detection - Frames Output",
-    description="Upload a short video to save detected frames as images and optionally generate a confidence score graph."
-)
+with gr.Blocks() as iface:
+    gr.Markdown("# YOLOv8 Object Detection - Real-time Frame Output")
+    gr.Markdown("Upload a short video to view frames with detections immediately in a gallery. Optionally generate a confidence score graph.")
+    
+    with gr.Row():
+        video_input = gr.Video(label="Upload Video")
+        output_folder = gr.Textbox(label="Output Folder", value="detected_frames")
+        plot_graphs = gr.Checkbox(label="Generate Confidence Score Graph", value=False)
+    
+    submit_button = gr.Button("Process Video")
+    
+    status_output = gr.Text(label="Status")
+    gallery_output = gr.Gallery(label="Detected Frames and Graph", preview=True, columns=3)
+    
+    submit_button.click(
+        fn=process_video,
+        inputs=[video_input, output_folder, plot_graphs],
+        outputs=[status_output, gallery_output],
+        concurrency_limit=1
+    )
 
 if __name__ == "__main__":
     iface.launch()