Browse Source

Upgrading the operator example to use `dora-arrow`

tags/v0.2.3-rc
haixuanTao 2 years ago
parent
commit
fa17544d8d
5 changed files with 49 additions and 10 deletions
  1. +14
    -2
      examples/python-operator-dataflow/no_webcam.py
  2. +14
    -3
      examples/python-operator-dataflow/object_detection.py
  3. +11
    -3
      examples/python-operator-dataflow/plot.py
  4. +1
    -0
      examples/python-operator-dataflow/requirements.txt
  5. +9
    -2
      examples/python-operator-dataflow/webcam.py

+ 14
- 2
examples/python-operator-dataflow/no_webcam.py View File

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


+ 14
- 3
examples/python-operator-dataflow/object_detection.py View File

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

+ 11
- 3
examples/python-operator-dataflow/plot.py View File

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


+ 1
- 0
examples/python-operator-dataflow/requirements.txt View File

@@ -43,3 +43,4 @@ thop>=0.1.1 # FLOPs computation
# roboflow

opencv-python>=4.1.1
pyarrow

+ 9
- 2
examples/python-operator-dataflow/webcam.py View File

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


Loading…
Cancel
Save