app.py

app.py

@@ -1,16 +1,12 @@
 import torch
-from yolov5.utils.general import non_max_suppression, scale_coords
-from yolov5.utils.plots import plot_boxes
-from yolov5.models.common import DetectMultiBackend
 import cv2
-from PIL import Image
 import numpy as np
 
 # Load your YOLO model
 model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt') 
 
 def detect(frame):
-    # Convert the frame to an image format YOLOv5 can process
+    # Perform inference on the current frame
     results = model(frame)
     return results
 
@@ -25,14 +21,18 @@ while cap.isOpened():
     # Perform inference on the current frame
     results = detect(frame)
 
-    # Display the frame with results (optional)
-    frame = np.array(results.render()[0])  # render predictions on the frame
-    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)  # Convert from RGB to BGR for OpenCV
-    cv2.imshow("Detection", frame)
+    # Render predictions on the frame (draw bounding boxes)
+    frame_with_boxes = np.array(results.render()[0])  # render predictions on the frame
+
+    # Convert from RGB to BGR for OpenCV (since OpenCV uses BGR)
+    frame_with_boxes = cv2.cvtColor(frame_with_boxes, cv2.COLOR_RGB2BGR)  
+
+    # Display the frame with results
+    cv2.imshow("Detection", frame_with_boxes)
 
     # Break loop if 'q' is pressed
     if cv2.waitKey(1) & 0xFF == ord('q'):
         break
 
 cap.release()
-cv2.destroyAllWindows()
+cv2.destroyAllWindows()