added label server

src/perception/lidar_object_detection/lidar_object_detection/config/nuscenes_config.yaml

@@ -3,5 +3,4 @@ lidar_object_detection_node:
     lidar_topic: /LIDAR_TOP
     model_config_path: /home/bolty/OpenPCDet/tools/cfgs/nuscenes_models/transfusion_lidar.yaml
     model_path: /home/bolty/OpenPCDet/models/cbgs_transfusion_lidar.pth
-  detection_options:
     enable_detection: true

src/perception/lidar_object_detection/lidar_object_detection/lidar_object_detection/label_server.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+
+import rclpy
+from rclpy.node import Node
+from vision_msgs.msg import VisionInfo
+
+class LabelServer(Node):
+    def __init__(self):
+        super().__init__('label_server')
+        self.publisher_ = self.create_publisher(VisionInfo, 'vision_info', 10)
+        self.label_mapping = {
+            1: "car",
+            2: "pedestrian",
+            3: "cyclist",
+        }
+        self.publish_labels()
+
+    def publish_labels(self):
+        vision_info_msg = VisionInfo()
+        vision_info_msg.database_location = "memory"
+        vision_info_msg.database_version = "1.0"
+        for label_id, class_name in self.label_mapping.items():
+            vision_info_msg.class_map[label_id] = class_name
+        self.publisher_.publish(vision_info_msg)
+
+def main(args=None):
+    rclpy.init(args=args)
+    label_server = LabelServer()
+    rclpy.spin(label_server)
+    label_server.destroy_node()
+    rclpy.shutdown()
+
+if __name__ == '__main__':
+    main()

src/perception/lidar_object_detection/lidar_object_detection/lidar_object_detection/lidar_object_detection_node.py

@@ -1,6 +1,6 @@
 # pylint: disable=wrong-import-position
-from sensor_msgs.msg import PointCloud2, PointField
-from vision_msgs.msg import ObjectHypothesisWithPose, Detection3D, Detection3DArray
+from sensor_msgs.msg import PointCloud2
+from vision_msgs.msg import ObjectHypothesisWithPose, Detection3D, Detection3DArray, VisionInfo
 from visualization_msgs.msg import Marker, MarkerArray
 from pcdet.config import cfg, cfg_from_yaml_file
 from pcdet.datasets import DatasetTemplate
@@ -14,7 +14,6 @@
 import sys
 sys.path.append("/home/bolty/OpenPCDet")
 
-
 class LidarObjectDetection(Node):
     def __init__(self):
         super().__init__('lidar_object_detection')
@@ -25,8 +24,14 @@ def __init__(self):
         self.model_config_path = self.get_parameter("model_config_path").value
         self.lidar_data = self.get_parameter("lidar_topic").value
         self.publish_detection = self.get_parameter(
-            'detection_options.enable_detection').get_parameter_value().bool_value
-
+            'enable_detection').get_parameter_value().bool_value
+
+        self.label_mapping = {}
+        self.subscription = self.create_subscription(
+            VisionInfo,
+            'vision_info',
+            self.vision_info_callback,
+            10)
         self.viz_publisher = self.create_publisher(MarkerArray, "/lidar_detections_viz", 10)
         self.detections_publisher = self.create_publisher(Detection3DArray, "/lidar_detections", 10)
 
@@ -50,6 +55,9 @@ def __init__(self):
         self.model.load_params_from_file(filename=args.ckpt, logger=self.logger, to_cpu=True)
         self.model.cuda()
         self.model.eval()
+
+    def vision_info_callback(self, msg):
+        self.label_mapping = msg.class_map
 
     def point_cloud_callback(self, msg):
         points = self.pointcloud2_to_xyz_array(msg)
@@ -107,6 +115,10 @@ def publish_bounding_boxes(self, pointcloud_msg, pred_dicts, original_timestamp)
                 marker.color.r = 1.0
                 marker.color.g = 0.0
                 marker.color.b = float(pred_dicts[0]["pred_labels"][idx]) / 3
+
+                class_id = pred_dicts[0]["pred_labels"][idx]
+                class_name = self.label_mapping.get(class_id, "unknown")
+                marker.text = class_name
                 marker_array.markers.append(marker)
 
                 detection = Detection3D()

src/perception/lidar_object_detection/lidar_object_detection/setup.py

@@ -24,7 +24,8 @@
     tests_require=['pytest'],
     entry_points={
         'console_scripts': [
-            'lidar_object_detection_node = lidar_object_detection.lidar_object_detection_node:main'
+            'lidar_object_detection_node = lidar_object_detection.lidar_object_detection_node:main',
+            'label_server = lidar_object_detection.label_server:main',
         ],
     },
 )