Replace yolov5 with yolov8

.gitignore

@@ -7,6 +7,7 @@
 
 # Remove arrow file from dora-record
 **/*.arrow
+*.pt
 
 # Removing images.
 *.jpg

examples/python-dataflow/object_detection.py

@@ -1,15 +1,14 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-import os
 import cv2
 import numpy as np
-import torch
+from ultralytics import YOLO
 
 from dora import Node
+import pyarrow as pa
 
-# Reload only if on Windows
-model = torch.hub.load("ultralytics/yolov5", "yolov5n")
+model = YOLO("yolov8n.pt")
 
 node = Node()
 
@@ -23,9 +22,14 @@ for event in node:
             frame = cv2.imdecode(frame, -1)
             frame = frame[:, :, ::-1]  # OpenCV image (BGR to RGB)
             results = model(frame)  # includes NMS
-            arrays = np.array(results.xyxy[0].cpu()).tobytes()
+            # Process results
+            boxes = np.array(results[0].boxes.xyxy.cpu())
+            conf = np.array(results[0].boxes.conf)
+            label = np.array(results[0].boxes.cls)
+            # concatenate them together
+            arrays = np.concatenate((boxes, conf[:, None], label[:, None]), axis=1)
 
-            node.send_output("bbox", arrays, event["metadata"])
+            node.send_output("bbox", pa.array(arrays.ravel()), event["metadata"])
         else:
             print("[object detection] ignoring unexpected input:", event_id)
     elif event_type == "STOP":

examples/python-operator-dataflow/object_detection.py

@@ -4,9 +4,9 @@
 
 import numpy as np
 import pyarrow as pa
-import torch
-import os
+
 from dora import DoraStatus
+from ultralytics import YOLO
 
 pa.array([])
 
@@ -20,7 +20,7 @@ class Operator:
     """
 
     def __init__(self):
-        self.model = torch.hub.load("ultralytics/yolov5", "yolov5n")
+        self.model = YOLO("yolov8n.pt")
 
     def on_event(
         self,
@@ -50,6 +50,12 @@ class Operator:
         frame = dora_input["value"].to_numpy().reshape((CAMERA_HEIGHT, CAMERA_WIDTH, 3))
         frame = frame[:, :, ::-1]  # OpenCV image (BGR to RGB)
         results = self.model(frame)  # includes NMS
-        arrays = pa.array(np.array(results.xyxy[0].cpu()).ravel())
+        # Process results
+        boxes = np.array(results[0].boxes.xyxy.cpu())
+        conf = np.array(results[0].boxes.conf)
+        label = np.array(results[0].boxes.cls)
+        # concatenate them together
+        arrays = np.concatenate((boxes, conf[:, None], label[:, None]), axis=1)
+
         send_output("bbox", arrays, dora_input["metadata"])
         return DoraStatus.CONTINUE