diff --git a/examples/python-dataflow/.gitignore b/examples/python-dataflow/.gitignore new file mode 100644 index 00000000..eede66d8 --- /dev/null +++ b/examples/python-dataflow/.gitignore @@ -0,0 +1 @@ +*.pt \ No newline at end of file diff --git a/examples/python-dataflow/dataflow.yml b/examples/python-dataflow/dataflow.yml index 8561be2c..26beefc0 100644 --- a/examples/python-dataflow/dataflow.yml +++ b/examples/python-dataflow/dataflow.yml @@ -3,18 +3,25 @@ communication: prefix: /example-python-dataflow nodes: - - id: node-1 - operators: - - id: python_webcam - python: webcam_operator.py + - id: webcam + operator: + python: webcam.py inputs: - timer: dora/timer/millis/100 + timer: dora/timer/millis/500 outputs: - image + + - id: object_detection + operator: + python: object_detection.py + inputs: + image: webcam/image + outputs: + - bbox - - id: node-2 - operators: - - id: python_plot - python: plot_cv2.py + - id: plot + operator: + python: plot.py inputs: - image: node-1/python_webcam/image + image: webcam/image + bbox: object_detection/bbox diff --git a/examples/python-dataflow/plot_cv2.py b/examples/python-dataflow/object_detection.py similarity index 60% rename from examples/python-dataflow/plot_cv2.py rename to examples/python-dataflow/object_detection.py index 50bc9e2b..5737fcfc 100644 --- a/examples/python-dataflow/plot_cv2.py +++ b/examples/python-dataflow/object_detection.py @@ -1,8 +1,8 @@ from enum import Enum from typing import Callable -import cv2 import numpy as np +import torch class DoraStatus(Enum): @@ -18,7 +18,7 @@ class Operator: """ def __init__(self): - self.counter = 0 + self.model = torch.hub.load("ultralytics/yolov5", "yolov5n") def on_input( self, @@ -33,17 +33,13 @@ class Operator: value (bytes): Bytes message of the input send_output (Callable[[str, bytes]]): Function enabling sending output back to dora. """ - self.counter += 1 - if input_id == "image": - frame = np.frombuffer(value, dtype="uint8") - frame = np.reshape(frame, (480, 640, 3)) - cv2.imshow("frame", frame) - if cv2.waitKey(1) & 0xFF == ord("q"): - return DoraStatus.STOP - if self.counter > 20: - return DoraStatus.STOP - else: - return DoraStatus.CONTINUE - - def drop_operator(self): - cv2.destroyAllWindows() + + frame = np.frombuffer(value, dtype="uint8") + frame = np.reshape(frame, (480, 640, 3))[ + :, :, ::-1 + ] # OpenCV image (BGR to RGB) + + results = self.model(frame) # includes NMS + arrays = np.array(results.xyxy[0].cpu()).tobytes() + send_output("bbox", arrays) + return DoraStatus.CONTINUE diff --git a/examples/python-dataflow/plot.py b/examples/python-dataflow/plot.py new file mode 100644 index 00000000..2ce8be30 --- /dev/null +++ b/examples/python-dataflow/plot.py @@ -0,0 +1,66 @@ +from enum import Enum +from typing import Callable + +import cv2 +import numpy as np + + +class DoraStatus(Enum): + CONTINUE = 0 + STOP = 1 + + +class Operator: + """ + Example operator incrementing a counter every times its been called. + + The current value of the counter is sent back to dora on `counter`. + """ + + def __init__(self): + self.image = [] + + def on_input( + self, + input_id: str, + value: bytes, + send_output: Callable[[str, bytes], None], + ) -> DoraStatus: + """Handle input by incrementing count by one. + + Args: + input_id (str): Id of the input declared in the yaml configuration + value (bytes): Bytes message of the input + send_output (Callable[[str, bytes]]): Function enabling sending output back to dora. + """ + if input_id == "image": + frame = np.frombuffer(value, dtype="uint8") + frame = np.reshape(frame, (480, 640, 3)) + self.image = frame + elif input_id == "bbox" and len(self.image) != 0: + bboxs = np.frombuffer(value, dtype="float32") + bboxs = np.reshape(bboxs, (-1, 6)) + for bbox in bboxs: + [ + min_x, + min_y, + max_x, + max_y, + _confidence, + _class_label, + ] = bbox + cv2.rectangle( + self.image, + (int(min_x), int(min_y)), + (int(max_x), int(max_y)), + (0, 255, 0), + 2, + ) + cv2.imshow("frame", self.image) + if cv2.waitKey(1) & 0xFF == ord("q"): + return DoraStatus.STOP + + return DoraStatus.CONTINUE + + def drop_operator(self): + cv2.destroyAllWindows() diff --git a/examples/python-dataflow/requirements.txt b/examples/python-dataflow/requirements.txt index fb10efbf..f3a6b71a 100644 --- a/examples/python-dataflow/requirements.txt +++ b/examples/python-dataflow/requirements.txt @@ -1 +1,15 @@ -pip install opencv-python \ No newline at end of file +# YOLOv5 requirements +# Usage: pip install -r requirements.txt + +# Base ---------------------------------------- +matplotlib>=3.2.2 +numpy>=1.18.5 +opencv-python>=4.1.1 +Pillow>=7.1.2 +PyYAML>=5.3.1 +requests>=2.23.0 +scipy>=1.4.1 +torch>=1.7.0 +torchvision>=0.8.1 +tqdm>=4.64.0 +protobuf<=3.20.1 # https://github.com/ultralytics/yolov5/issues/8012 diff --git a/examples/python-dataflow/webcam_operator.py b/examples/python-dataflow/webcam.py similarity index 90% rename from examples/python-dataflow/webcam_operator.py rename to examples/python-dataflow/webcam.py index 8eedc550..ab704c13 100644 --- a/examples/python-dataflow/webcam_operator.py +++ b/examples/python-dataflow/webcam.py @@ -18,6 +18,7 @@ class Operator: def __init__(self): self.video_capture = cv2.VideoCapture(0) + self.counter = 0 def on_input( self, @@ -39,6 +40,10 @@ class Operator: else: print("did not sent video") + self.counter += 1 + if self.counter > 100: + return DoraStatus.STOP + return DoraStatus.CONTINUE def drop_operator(self):