App.py

app.py

@@ -2,16 +2,37 @@ 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(image):
-    results = model(image)
-    return results 
+def detect(frame):
+    # Convert the frame to an image format YOLOv5 can process
+    results = model(frame)
+    return results
 
-# Example inference
-image = Image.open("test.jpg")
-results = detect(image)
-print(results)
+# Open the video file (or camera)
+cap = cv2.VideoCapture("test_video.mp4")  # replace with your video path or camera index (0 for webcam)
+
+while cap.isOpened():
+    ret, frame = cap.read()
+    if not ret:
+        break
+    
+    # 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)
+
+    # Break loop if 'q' is pressed
+    if cv2.waitKey(1) & 0xFF == ord('q'):
+        break
+
+cap.release()
+cv2.destroyAllWindows()