190 - Add y distance of traffic light. (#208)

fix: traffic light detection + add y distance of traffic light.

code/perception/src/traffic_light_node.py

@@ -8,6 +8,7 @@
 from rospy.numpy_msg import numpy_msg
 from sensor_msgs.msg import Image as ImageMsg
 from perception.msg import TrafficLightState
+from std_msgs.msg import Int16
 from cv_bridge import CvBridge
 from traffic_light_detection.src.traffic_light_detection.traffic_light_inference \
     import TrafficLightInference  # noqa: E501
@@ -45,6 +46,12 @@ def setup_traffic_light_publishers(self):
             topic=f"/paf/{self.role_name}/{self.side}/traffic_light_state",
             qos_profile=1
         )
+        self.traffic_light_distance_publisher = self.new_publisher(
+            msg_type=Int16,
+            topic=f"/paf/{self.role_name}/{self.side}" +
+            "/traffic_light_y_distance",
+            qos_profile=1
+        )
 
     def auto_invalidate_state(self):
         while True:
@@ -57,9 +64,12 @@ def auto_invalidate_state(self):
                 msg = TrafficLightState()
                 msg.state = 0
                 self.traffic_light_publisher.publish(msg)
+                self.traffic_light_distance_publisher.publish(Int16(0))
                 self.last_info_time = None
 
     def handle_camera_image(self, image):
+        distance = int(image.header.frame_id)
+
         cv2_image = self.bridge.imgmsg_to_cv2(image)
         rgb_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)
         result, data = self.classifier(cv2_image)
@@ -77,6 +87,9 @@ def handle_camera_image(self, image):
             msg = TrafficLightState()
             msg.state = state
             self.traffic_light_publisher.publish(msg)
+
+            if distance is not None:
+                self.traffic_light_distance_publisher.publish(Int16(distance))
         else:
             self.last_state = state
 

code/perception/src/vision_node.py

@@ -19,6 +19,7 @@
 from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks
 import numpy as np
 from ultralytics import NAS, YOLO, RTDETR, SAM, FastSAM
+import asyncio
 import rospy
 
 """
@@ -394,7 +395,9 @@ def predict_ultralytics(self, image):
                                                 dtype=torch.uint8)
 
         if 9 in output[0].boxes.cls:
-            self.process_traffic_lights(output[0], cv_image, image.header)
+            asyncio.run(self.process_traffic_lights(output[0],
+                                                    cv_image,
+                                                    image.header))
 
         c_boxes = torch.stack(c_boxes)
         box = draw_bounding_boxes(image_np_with_detections,
@@ -433,9 +436,12 @@ def process_traffic_lights(self, prediction, cv_image, image_header):
             box = box[0:4].astype(int)
             segmented = cv_image[box[1]:box[3], box[0]:box[2]]
 
+            traffic_light_y_distance = box[1]
+
             traffic_light_image = self.bridge.cv2_to_imgmsg(segmented,
                                                             encoding="rgb8")
             traffic_light_image.header = image_header
+            traffic_light_image.header.frame_id = str(traffic_light_y_distance)
             self.traffic_light_publisher.publish(traffic_light_image)
 
         # locals().clear()