Update app.py

app.py

@@ -1,10 +1,13 @@
 import cv2
 import torch
 import gradio as gr
+import numpy as np
 from ultralytics import YOLO
 
-# Load YOLOv8 model
+# Load YOLOv8 model and set device (GPU if available)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = YOLO('./data/best.pt')  # Path to your model
+model.to(device)
 
 # Define the function that processes the uploaded video
 def process_video(video):
@@ -16,10 +19,9 @@ def process_video(video):
     frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = input_video.get(cv2.CAP_PROP_FPS)
 
-    # Define output video writer
-    output_video_path = "/mnt/data/output_video.mp4"  # Path to save the output video
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    output_video = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
+    # Resize to reduce computation (optional)
+    new_width, new_height = 640, 480  # Resize to 640x480 resolution
+    frame_width, frame_height = new_width, new_height
 
     while True:
         # Read a frame from the video
@@ -27,27 +29,39 @@ def process_video(video):
         if not ret:
             break  # End of video
 
+        # Resize the frame to reduce computational load
+        frame = cv2.resize(frame, (new_width, new_height))
+
         # Perform inference on the frame
-        results = model(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
+            # Annotate the frame with bounding boxes
+            annotated_frame = results[0].plot()  # Plot the frame with bounding boxes
+
+            # Convert the annotated frame to RGB format for displaying
+            annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
 
-        # The results object contains annotations for the frame
-        annotated_frame = results[0].plot()  # Plot the frame with bounding boxes
+            # Display the frame with detections
+            cv2.imshow("Detected Frame", annotated_frame_rgb)
 
-        # Write the annotated frame to the output video
-        output_video.write(annotated_frame)
+            # Wait for a key press (optional: press 'q' to quit early)
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
 
     # Release resources
     input_video.release()
-    output_video.release()
+    cv2.destroyAllWindows()
 
-    return output_video_path
+    return "Video processing complete!"
 
 # Create a Gradio interface for video upload
 iface = gr.Interface(fn=process_video, 
                      inputs=gr.Video(label="Upload Video"),  # Updated line
-                     outputs="file", 
-                     title="YOLOv8 Object Detection on Video",
-                     description="Upload a video for object detection using YOLOv8")
+                     outputs=gr.Textbox(label="Processing Status"),  # Output text showing processing status
+                     title="YOLOv8 Object Detection - Real-Time Display",
+                     description="Upload a video for object detection using YOLOv8. The frames with detections will be shown in real-time.")
 
 # Launch the interface
 iface.launch()