adding object detection to python example

examples/python-dataflow/.gitignore

@@ -0,0 +1 @@
+*.pt

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

examples/python-dataflow/plot_cv2.py → 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

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()

examples/python-dataflow/requirements.txt

@@ -1 +1,15 @@
-pip install opencv-python
+# 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

examples/python-dataflow/webcam_operator.py → 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):