From fa17544d8d3eac2c8a3c1a4e76137aafcfcc702b Mon Sep 17 00:00:00 2001 From: haixuanTao Date: Mon, 17 Apr 2023 20:59:17 +0800 Subject: [PATCH] Upgrading the operator example to use `dora-arrow` --- examples/python-operator-dataflow/no_webcam.py | 16 ++++++++++++++-- .../object_detection.py | 17 ++++++++++++++--- examples/python-operator-dataflow/plot.py | 14 +++++++++++--- .../python-operator-dataflow/requirements.txt | 1 + examples/python-operator-dataflow/webcam.py | 11 +++++++++-- 5 files changed, 49 insertions(+), 10 deletions(-) diff --git a/examples/python-operator-dataflow/no_webcam.py b/examples/python-operator-dataflow/no_webcam.py index 50023cce..60a4f950 100755 --- a/examples/python-operator-dataflow/no_webcam.py +++ b/examples/python-operator-dataflow/no_webcam.py @@ -4,16 +4,28 @@ import time import urllib.request +import cv2 import numpy as np +import pyarrow as pa from dora import Node print("Hello from no_webcam.py") -req = urllib.request.urlopen("https://ultralytics.com/images/zidane.jpg") +CAMERA_WIDTH = 640 +CAMERA_HEIGHT = 480 +# Preprocessing the image +req = urllib.request.urlopen( + "https://img0.baidu.com/it/u=2940037857,1417768899&fm=253&fmt=auto&app=138&f=PNG?w=724&h=500" +) # This image works in china better arr = np.asarray(bytearray(req.read()), dtype=np.uint8) +image = cv2.imdecode(arr, -1)[:, :, :3] +image = cv2.resize(image, (CAMERA_WIDTH, CAMERA_HEIGHT)) + +# Numpy -> Arrow +image = pa.array(image.flatten().view(np.uint8)) node = Node() start = time.time() @@ -24,7 +36,7 @@ while time.time() - start < 20: match event["type"]: case "INPUT": print("received input", event["id"]) - node.send_output("image", arr.tobytes()) + node.send_output("image", image) case "STOP": print("received stop") case other: diff --git a/examples/python-operator-dataflow/object_detection.py b/examples/python-operator-dataflow/object_detection.py index d06974cc..e2c5c8cc 100755 --- a/examples/python-operator-dataflow/object_detection.py +++ b/examples/python-operator-dataflow/object_detection.py @@ -5,10 +5,16 @@ from typing import Callable import cv2 import numpy as np +import pyarrow as pa import torch from dora import DoraStatus +pa.array([]) + +CAMERA_WIDTH = 640 +CAMERA_HEIGHT = 480 + class Operator: """ @@ -38,10 +44,15 @@ class Operator: send_output (Callable[[str, bytes]]): Function enabling sending output back to dora. """ - frame = np.frombuffer(dora_input["data"], dtype="uint8") - frame = cv2.imdecode(frame, -1) + 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 = np.array(results.xyxy[0].cpu()).tobytes() + arrays = pa.array( + np.array(results.xyxy[0].cpu()).flatten().view(np.uint8) + ) send_output("bbox", arrays, dora_input["metadata"]) return DoraStatus.CONTINUE diff --git a/examples/python-operator-dataflow/plot.py b/examples/python-operator-dataflow/plot.py index 8c50db1d..92f4cccc 100755 --- a/examples/python-operator-dataflow/plot.py +++ b/examples/python-operator-dataflow/plot.py @@ -6,11 +6,16 @@ from typing import Callable import cv2 import numpy as np +import pyarrow as pa from utils import LABELS from dora import DoraStatus +pa.array([]) + CI = os.environ.get("CI") +CAMERA_WIDTH = 640 +CAMERA_HEIGHT = 480 font = cv2.FONT_HERSHEY_SIMPLEX @@ -49,15 +54,18 @@ class Operator: send_output (Callable[[str, bytes]]): Function enabling sending output back to dora. """ if dora_input["id"] == "image": - frame = np.frombuffer(dora_input["data"], dtype="uint8") - frame = cv2.imdecode(frame, -1) + frame = ( + dora_input["value"] + .to_numpy() + .reshape((CAMERA_HEIGHT, CAMERA_WIDTH, 3)) + ) self.image = frame self.image_messages += 1 print("received " + str(self.image_messages) + " images") elif dora_input["id"] == "bbox" and len(self.image) != 0: - bboxs = np.frombuffer(dora_input["data"], dtype="float32") + bboxs = dora_input["value"].to_numpy().view(np.float32) self.bboxs = np.reshape(bboxs, (-1, 6)) self.bounding_box_messages += 1 diff --git a/examples/python-operator-dataflow/requirements.txt b/examples/python-operator-dataflow/requirements.txt index 55f71178..825cd671 100644 --- a/examples/python-operator-dataflow/requirements.txt +++ b/examples/python-operator-dataflow/requirements.txt @@ -43,3 +43,4 @@ thop>=0.1.1 # FLOPs computation # roboflow opencv-python>=4.1.1 +pyarrow \ No newline at end of file diff --git a/examples/python-operator-dataflow/webcam.py b/examples/python-operator-dataflow/webcam.py index e2c73ab4..f5f36019 100755 --- a/examples/python-operator-dataflow/webcam.py +++ b/examples/python-operator-dataflow/webcam.py @@ -5,9 +5,13 @@ import time from typing import Callable import cv2 +import pyarrow as pa from dora import DoraStatus +CAMERA_WIDTH = 640 +CAMERA_HEIGHT = 480 + class Operator: """ @@ -17,19 +21,22 @@ class Operator: def __init__(self): self.video_capture = cv2.VideoCapture(0) self.start_time = time.time() + self.video_capture.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH) + self.video_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT) def on_event( self, - dora_event: dict, + dora_event: str, send_output: Callable[[str, bytes], None], ) -> DoraStatus: match dora_event["type"]: case "INPUT": ret, frame = self.video_capture.read() + frame = cv2.resize(frame, (CAMERA_WIDTH, CAMERA_HEIGHT)) if ret: send_output( "image", - cv2.imencode(".jpg", frame)[1].tobytes(), + pa.array(frame.flatten()), dora_event["metadata"], ) case "STOP":