From d7ab431fd6d1b4f44da36c9674bbfdea9c8eb71c Mon Sep 17 00:00:00 2001
From: Raphael Maenle <rm@d-aria.at>
Date: Mon, 14 Mar 2022 17:42:34 +0100
Subject: [PATCH] publishes bounding boxes

- bounding boxes found in image over 0.1
  score are published as a Detection2D
  array
- post-processing moved to seperate thread
---
 image_publisher.py |  2 +-
 ros_inference.py   | 28 +++++++++++++++++++++-------
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/image_publisher.py b/image_publisher.py
index 043d410..0377a24 100644
--- a/image_publisher.py
+++ b/image_publisher.py
@@ -10,7 +10,7 @@ class MinimalPublisher(Node):
       def __init__(self):
          super().__init__('minimal_publisher')
          self.publisher_ = self.create_publisher(Image, '/camera/color/image_raw', 10)
-         timer_period = 0.5  # seconds
+         timer_period = 1/30  # seconds
          self.timer = self.create_timer(timer_period, self.timer_callback)
          self.i = 0
          self.im_list = []
diff --git a/ros_inference.py b/ros_inference.py
index fade56d..bbd10b2 100644
--- a/ros_inference.py
+++ b/ros_inference.py
@@ -24,8 +24,9 @@ from tensorflow.image import combined_non_max_suppression
 import rclpy
 from rclpy.node import Node
 
-from std_msgs.msg import String as StringMsg
 from sensor_msgs.msg import Image as ImageMsg
+from vision_msgs.msg import Detection2DArray, Detection2D, BoundingBox2D, ObjectHypothesisWithPose
+from geometry_msgs.msg import Pose2D
 
 from cv_bridge import CvBridge
 
@@ -369,7 +370,7 @@ class HailoNode(Node):
     def __init__(self):
         super().__init__('hailo_image_subscriber')
         self.sub = self.create_subscription(ImageMsg, '/camera/color/image_raw', self.image_callback, 10)
-        self.pub = self.create_publisher(ImageMsg, '/hailo_bounding_boxes', 10)
+        self.pub = self.create_publisher(Detection2DArray, '/hailo_bounding_boxes', 10)
 
         self.bridge = CvBridge()
 
@@ -383,7 +384,6 @@ class HailoNode(Node):
         self.hailo.start_hailo_thread()
 
     def image_callback(self, data):
-        print("Received an image!")
         img = self.convert(data)
         self.image_infer(img)
 
@@ -396,19 +396,33 @@ class HailoNode(Node):
         while(out == None):
             time.sleep(0.0001)
             out = self.hailo.hailo_output()
+
+        Thread(target=self._thread_postprocessing, args=[out]).start()
         
+
+    def _thread_postprocessing(self, out):
         logits = self.processor.postprocessing(out)
 
         labels = get_labels("data/daria_labels.json")
+
+        detection_array = Detection2DArray()
+
         for bb in range(len(logits['detection_boxes'].numpy()[0])):
             boxes = logits['detection_boxes'].numpy()[0][bb]
             classes = logits['detection_classes'][0][bb]
             scores = logits['detection_scores'].numpy()[0][bb]
             if(scores > 0.01):
-                print(boxes)
-                print(labels[classes])
-                print(scores)
-                print("-----")
+                bb = BoundingBox2D()
+                bb.center = Pose2D()
+                bb.center.x = float(boxes[0])
+                bb.center.y = float(boxes[1])
+                bb.size_x = float(boxes[2])
+                bb.size_y = float(boxes[3])
+                oh = ObjectHypothesisWithPose(id=str(labels[classes]), score=float(scores))
+                detection = Detection2D(results=[oh], bbox=bb)
+                detection_array.detections.append(detection)
+
+        self.pub.publish(detection_array)
 
     # convert ros image to PIL image
     def convert(self, ros_image):