publish empty detection in 3d node

mgonzs13 · Oct 20, 2023 · ab4216a · ab4216a
1 parent bbd7cb9
commit ab4216a
Showing 1 changed file with 25 additions and 14 deletions.
diff --git a/yolov8_ros/yolov8_ros/detect_3d_node.py b/yolov8_ros/yolov8_ros/detect_3d_node.py
@@ -15,7 +15,7 @@
 
 
 import numpy as np
-from typing import Tuple
+from typing import Tuple, List
 
 import rclpy
 from rclpy.node import Node
@@ -66,45 +66,56 @@ def __init__(self) -> None:
             (self.points_sub, self.detections_sub), 10, 0.5)
         self._synchronizer.registerCallback(self.on_detections)
 
-    def on_detections(self,
-                      points_msg: PointCloud2,
-                      detections_msg: DetectionArray,
-                      ) -> None:
+    def on_detections(
+        self,
+        points_msg: PointCloud2,
+        detections_msg: DetectionArray,
+    ) -> None:
+
+        new_detections_msg = DetectionArray()
+        new_detections_msg.header = detections_msg.header
+        new_detections_msg.detections = self.process_detections(
+            points_msg, detections_msg)
+        self._pub.publish(new_detections_msg)
+
+    def process_detections(
+            self,
+            points_msg: PointCloud2,
+            detections_msg: DetectionArray
+    ) -> List[Detection]:
 
         # check if there are detections
         if not detections_msg.detections:
-            return
+            return []
 
         transform = self.get_transform(points_msg.header.frame_id)
 
         if transform is None:
-            return
+            return []
 
+        new_detections = []
         points = np.frombuffer(points_msg.data, np.float32).reshape(
             points_msg.height, points_msg.width, -1)[:, :, :3]
 
-        new_detections_msg = DetectionArray()
-        new_detections_msg.header = detections_msg.header
-
         for detection in detections_msg.detections:
             bbox3d = self.convert_bb_to_3d(points, detection)
 
             if bbox3d is not None:
-                new_detections_msg.detections.append(detection)
+                new_detections.append(detection)
 
                 bbox3d = Detect3DNode.transform_3d_box(
                     bbox3d, transform[0], transform[1])
                 bbox3d.frame_id = self.target_frame
-                new_detections_msg.detections[-1].bbox3d = bbox3d
+                new_detections[-1].bbox3d = bbox3d
 
                 keypoints3d = self.convert_keypoints_to_3d(
                     points, detection)
                 keypoints3d = Detect3DNode.transform_3d_keypoints(
                     keypoints3d, transform[0], transform[1])
                 keypoints3d.frame_id = self.target_frame
-                new_detections_msg.detections[-1].keypoints3d = keypoints3d
+                new_detections[-1].keypoints3d = keypoints3d
 
-        self._pub.publish(new_detections_msg)
+        return new_detections
 
     def convert_bb_to_3d(self,
                          points: np.ndarray,