Update app.py

app.py

@@ -2,82 +2,92 @@ import cv2
 import torch
 import gradio as gr
 import numpy as np
-from ultralytics import YOLO
+import os
 import matplotlib.pyplot as plt
+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}")
+print(f"CUDA available: {torch.cuda.is_available()}")
 
 # Load YOLOv8 model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = YOLO('./data/best.pt')  # Path to your model
-model.to(device)
-
-# List to store frames with detections
-frames_with_detections = []
-
-# Define the function to process the video
-def process_video(video):
-    # Open the video file
-    input_video = cv2.VideoCapture(video)  
-    frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = input_video.get(cv2.CAP_PROP_FPS)
-
-    # Resize frames to 640x480 (optional, to reduce computational load)
-    new_width, new_height = 640, 480
-
+print(f"Using device: {device}")
+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"
+    
+    # Create output folder if it doesn't exist
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+    
+    cap = cv2.VideoCapture(video)
+    if not cap.isOpened():
+        return "Error: Could not open video file"
+    
+    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
+    
     while True:
-        # Read a frame from the video
-        ret, frame = input_video.read()
-        if not ret:
-            break  # End of video
-
-        # Resize the frame
-        frame = cv2.resize(frame, (new_width, new_height))
-
-        # Perform inference on the frame
-        results = model(frame)  # Automatically uses GPU if available
-
-        # If there are detections
-        if len(results[0].boxes) > 0: 
-            boxes = results[0].boxes.xyxy.cpu().numpy()  # Get the bounding boxes
-
-            # Annotate the frame with bounding boxes
-            annotated_frame = results[0].plot()  
-
-            # Convert the frame to RGB
-            annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
-
-            # Append the frame with detection to list
-            frames_with_detections.append(annotated_frame_rgb)
-
-            # Create a simple bar chart to show the count of detected objects
-            fig, ax = plt.subplots()
-            ax.bar([1], [len(boxes)], color='blue')  # Bar for the current frame detection
-            ax.set_xlabel('Frame')
-            ax.set_ylabel('Number of Detections')
-            ax.set_title('Detection Count per Frame')
-
-            # Convert plot to an image to return it in Gradio output
-            plt.tight_layout()
-            plt.close(fig)
-
-            # Save the plot as an image in memory
-            buf = np.frombuffer(fig.canvas.print_to_buffer()[0], dtype=np.uint8)
-            img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
+        ret, frame = cap.read()
+        if not ret or frame_count > max_frames:
+            break
             
-            # Yield the detected frame and the graph at the same time
-            yield annotated_frame_rgb, img
-
-    # Release resources
-    input_video.release()
+        frame_count += 1
+        if frame_count % frame_skip != 0:
+            continue
+            
+        frame = cv2.resize(frame, (frame_width, frame_height))
+        print(f"Processing frame {frame_count}")
+        
+        # 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:
+            confs = results[0].boxes.conf.cpu().numpy()
+            confidence_scores.extend(confs)
+    
+    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)
+        plt.title('Distribution of Confidence Scores')
+        plt.xlabel('Confidence Score')
+        plt.ylabel('Frequency')
+        graph_path = os.path.join(output_folder, "confidence_histogram.png")
+        plt.savefig(graph_path)
+        plt.close()
+    
+    return f"Frames saved in {output_folder}. {f'Graph saved as {graph_path}' if graph_path else ''}"
 
 # Gradio interface
-with gr.Blocks() as demo:
-    with gr.Row():
-        video_input = gr.Video(label="Upload Video")
-        gallery_output = gr.Gallery(label="Detection Album")  # Removed style() method
-        graph_output = gr.Image(label="Detection Counts Graph", type="numpy")  # For displaying graph
-
-    video_input.change(process_video, inputs=video_input, outputs=[gallery_output, graph_output])
-
-# Launch the interface
-demo.launch()
+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."
+)
+
+if __name__ == "__main__":
+    iface.launch()