diff --git a/.gitignore b/.gitignore index 557fa2e1..59491033 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ # Remove arrow file from dora-record **/*.arrow +*.pt # Removing images. *.jpg diff --git a/examples/python-dataflow/object_detection.py b/examples/python-dataflow/object_detection.py index b9fff333..5b4788d8 100755 --- a/examples/python-dataflow/object_detection.py +++ b/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": diff --git a/examples/python-operator-dataflow/object_detection.py b/examples/python-operator-dataflow/object_detection.py index 90dbf449..d1aed848 100755 --- a/examples/python-operator-dataflow/object_detection.py +++ b/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