From 7b5c31d21400c180dd8dad7e377e4365812a9546 Mon Sep 17 00:00:00 2001
From: haixuanTao <tao.xavier@outlook.com>
Date: Sat, 20 Apr 2024 13:02:15 +0200
Subject: [PATCH] Update example to add boundingbox 2d

---
 examples/rerun-viewer/dataflow.yml        | 17 ++++++++-
 examples/rerun-viewer/object_detection.py | 45 +++++++++++++++++++++++
 examples/rerun-viewer/plot.py             | 12 +++---
 examples/rerun-viewer/webcam.py           |  5 ++-
 4 files changed, 70 insertions(+), 9 deletions(-)
 create mode 100755 examples/rerun-viewer/object_detection.py

diff --git a/examples/rerun-viewer/dataflow.yml b/examples/rerun-viewer/dataflow.yml
index 1b9a1c40e..5e179f027 100644
--- a/examples/rerun-viewer/dataflow.yml
+++ b/examples/rerun-viewer/dataflow.yml
@@ -14,15 +14,27 @@ nodes:
         IMAGE_HEIGHT: 540
 
 
+  - id: object_detection
+    custom:
+      source: ./object_detection.py
+      inputs:
+        image: webcam/image
+      outputs:
+        - bbox
+      envs:
+        IMAGE_WIDTH: 960
+        IMAGE_HEIGHT: 540
+
   - id: rerun
     custom:
       source: dora-rerun
       inputs:
         image: webcam/image
         text: webcam/text
+        boxes2d: object_detection/bbox
       envs:
-        IMAGE_WIDTH: 540
-        IMAGE_HEIGHT: 960
+        IMAGE_WIDTH: 960
+        IMAGE_HEIGHT: 540
         IMAGE_DEPTH: 3
 
   - id: matplotlib
@@ -30,6 +42,7 @@ nodes:
       source: ./plot.py
       inputs:
         image: webcam/image
+        bbox: object_detection/bbox
       envs:
         IMAGE_WIDTH: 960
         IMAGE_HEIGHT: 540
\ No newline at end of file
diff --git a/examples/rerun-viewer/object_detection.py b/examples/rerun-viewer/object_detection.py
new file mode 100755
index 000000000..2c606be98
--- /dev/null
+++ b/examples/rerun-viewer/object_detection.py
@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import os
+import cv2
+import numpy as np
+from ultralytics import YOLO
+
+from dora import Node
+import pyarrow as pa
+
+model = YOLO("yolov8n.pt")
+
+node = Node()
+
+IMAGE_WIDTH = int(os.getenv("IMAGE_WIDTH", 960))
+IMAGE_HEIGHT = int(os.getenv("IMAGE_HEIGHT", 540))
+
+for event in node:
+    event_type = event["type"]
+    if event_type == "INPUT":
+        event_id = event["id"]
+        if event_id == "image":
+            print("[object detection] received image input")
+            image = event["value"].to_numpy().reshape((IMAGE_HEIGHT, IMAGE_WIDTH, 3))
+
+            frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            frame = frame[:, :, ::-1]  # OpenCV image (BGR to RGB)
+            results = model(frame)  # includes NMS
+            # Process results
+            boxes = np.array(results[0].boxes.xywh.cpu())
+            conf = np.array(results[0].boxes.conf.cpu())
+            label = np.array(results[0].boxes.cls.cpu())
+            # concatenate them together
+            arrays = np.concatenate((boxes, conf[:, None], label[:, None]), axis=1)
+
+            node.send_output("bbox", pa.array(arrays.ravel()), event["metadata"])
+        else:
+            print("[object detection] ignoring unexpected input:", event_id)
+    elif event_type == "STOP":
+        print("[object detection] received stop")
+    elif event_type == "ERROR":
+        print("[object detection] error: ", event["error"])
+    else:
+        print("[object detection] received unexpected event:", event_type)
diff --git a/examples/rerun-viewer/plot.py b/examples/rerun-viewer/plot.py
index 649fc82c9..d6ec83893 100755
--- a/examples/rerun-viewer/plot.py
+++ b/examples/rerun-viewer/plot.py
@@ -49,17 +49,17 @@ def on_input(
             self.bboxs = np.reshape(bboxs, (-1, 6))
         for bbox in self.bboxs:
             [
-                min_x,
-                min_y,
-                max_x,
-                max_y,
+                x,
+                y,
+                w,
+                h,
                 confidence,
                 label,
             ] = bbox
             cv2.rectangle(
                 self.image,
-                (int(min_x), int(min_y)),
-                (int(max_x), int(max_y)),
+                (int(x - w / 2), int(y - h / 2)),
+                (int(x + w / 2), int(y + h / 2)),
                 (0, 255, 0),
                 2,
             )
diff --git a/examples/rerun-viewer/webcam.py b/examples/rerun-viewer/webcam.py
index f85901503..33a7950df 100755
--- a/examples/rerun-viewer/webcam.py
+++ b/examples/rerun-viewer/webcam.py
@@ -22,9 +22,12 @@
 start = time.time()
 
 # Run for 20 seconds
-while time.time() - start < 60:
+while time.time() - start < 1000:
     # Wait next dora_input
     event = node.next()
+    if event is None:
+        break
+
     event_type = event["type"]
     if event_type == "INPUT":
         ret, frame = video_capture.read()