Merge pull request #251 from dora-rs/add-arrow-example

Upgrading the operator example to use `dora-arrow`

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 @@
     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:

examples/python-operator-dataflow/object_detection.py

@@ -5,10 +5,16 @@
 
 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 @@ def on_input(
             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()).ravel().view(np.uint8)
+        )
         send_output("bbox", arrays, dora_input["metadata"])
         return DoraStatus.CONTINUE

examples/python-operator-dataflow/plot.py

@@ -6,11 +6,16 @@
 
 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 @@ def on_input(
             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

examples/python-operator-dataflow/requirements.txt

@@ -43,3 +43,4 @@ thop>=0.1.1  # FLOPs computation
 # roboflow
 
 opencv-python>=4.1.1
+pyarrow

examples/python-operator-dataflow/webcam.py

@@ -5,9 +5,13 @@
 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.ravel()),
                         dora_event["metadata"],
                     )
             case "STOP":