Update app.py

app.py

@@ -3,32 +3,26 @@ import torch
 import gradio as gr
 import numpy as np
 from ultralytics import YOLO
+import matplotlib.pyplot as plt
 
 # Load YOLOv8 model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = YOLO('./data/best.pt')  # Path to your model
 model.to(device)
 
-# Store frames with detected objects
+# List to store frames with detections
 frames_with_detections = []
-detection_counts = []
 
-# Define the function that processes the uploaded video
+# Define the function to process the video
 def process_video(video):
-    # video is now the file path string, not a file object
-    input_video = cv2.VideoCapture(video)  # Directly pass the path to cv2.VideoCapture
-
-    # Get frame width, height, and fps from input 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 to reduce computation (optional)
-    new_width, new_height = 640, 480  # Resize to 640x480 resolution
-    frame_width, frame_height = new_width, new_height
-
-    # Track detected objects by their bounding box coordinates
-    detected_boxes = set()
+    # Resize frames to 640x480 (optional, to reduce computational load)
+    new_width, new_height = 640, 480
 
     while True:
         # Read a frame from the video
@@ -36,56 +30,54 @@ def process_video(video):
         if not ret:
             break  # End of video
 
-        # Resize the frame to reduce computational load
+        # Resize the frame
         frame = cv2.resize(frame, (new_width, new_height))
 
         # Perform inference on the frame
         results = model(frame)  # Automatically uses GPU if available
 
-        # Check if any object was detected
-        if len(results[0].boxes) > 0:  # If there are detected objects
-            # Get the bounding boxes for each detected object
-            boxes = results[0].boxes.xyxy.cpu().numpy()  # Get xyxy coordinates
+        # If there are detections
+        if len(results[0].boxes) > 0: 
+            boxes = results[0].boxes.xyxy.cpu().numpy()  # Get the bounding boxes
 
-            # Loop through each detection and only show the frame for new objects
-            for box in boxes:
-                x1, y1, x2, y2 = box
-                detection_box = (x1, y1, x2, y2)
+            # Annotate the frame with bounding boxes
+            annotated_frame = results[0].plot()  
 
-                # Check if this box was already processed
-                if detection_box not in detected_boxes:
-                    # Add the box to the set to avoid repeating the detection
-                    detected_boxes.add(detection_box)
+            # Convert the frame to RGB
+            annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
 
-                    # Annotate the frame with bounding boxes
-                    annotated_frame = results[0].plot()  # Plot the frame with bounding boxes
+            # Append the frame with detection to list
+            frames_with_detections.append(annotated_frame_rgb)
 
-                    # Convert the annotated frame to RGB format for displaying
-                    annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
+            # 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')
 
-                    # Add this frame to the list of frames with detections
-                    frames_with_detections.append(annotated_frame_rgb)
+            # Convert plot to an image to return it in Gradio output
+            plt.tight_layout()
+            plt.close(fig)
 
-                    # Yield the latest frame immediately for Gradio's real-time display
-                    yield annotated_frame_rgb
+            # 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)
+            
+            # Yield the detected frame and the graph at the same time
+            yield annotated_frame_rgb, img
 
     # Release resources
     input_video.release()
 
-# Create a Gradio Blocks interface
+# Gradio interface
 with gr.Blocks() as demo:
-    # Define a file input for video upload
-    video_input = gr.Video(label="Upload Video")
-    
-    # Define the output area to show processed frames (gallery for continuous update)
-    gallery_output = gr.Gallery(label="Detection Album", show_label=True, columns=3)  # Display images in a row (album)
-
-    # Define the function to update frames in the album
-    def update_gallery(video):
-        return process_video(video)  # Return frames one by one as they are detected
-
-    # Connect the video input to the gallery update
-    video_input.change(update_gallery, inputs=video_input, outputs=gallery_output)
+    with gr.Row():
+        video_input = gr.Video(label="Upload Video")
+        gallery_output = gr.Gallery(label="Detection Album").style(columns=3)  # Display images in a row
+        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()