110 - Implement traffic light detection (#148)

* Add publisher

* Publish segmented traffic lights

* Implement TrafficLightNode

* Add TrafficLightState msg. WIP

* Add publisher

* Publish segmented traffic lights

* Implement TrafficLightNode

* Add TrafficLightState msg. WIP

* Added side view classification

* Finish traffic light node

* Add manual control launch file

* Make linter happy

* Add documentation

* Add missing traffic light detection model

* Fix color issues in rviz

* Limit simulator's max. RAM usage to prevent system crash

* fix: Linter fixes for other team's code

.gitignore

@@ -11,3 +11,5 @@ code/output
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
+
+*.tsv

build/docker-compose.yml

@@ -26,6 +26,10 @@ services:
     # image: carlasim/carla:0.9.14
     image: ghcr.io/una-auxme/paf23:leaderboard-2.0
     init: true
+    deploy:
+      resources:
+        limits:
+          memory: 16G
     expose:
       - 2000
       - 2001

code/agent/launch/agent_manual.launch

@@ -0,0 +1,31 @@
+<launch>
+  <arg name="role_name" default="hero" />
+  <arg name="control_loop_rate" default="0.05" />
+
+  <!-- perception -->
+  <include file="$(find perception)/launch/perception.launch">
+  </include>
+
+  <!--
+  <include file="$(find planning_runner)/launch/planning_runner.launch">
+  </include>
+
+  <include file="$(find mock)/launch/mock.launch">
+      <arg name="control_loop_rate" value="$(arg control_loop_rate)"/>
+      <arg name="role_name" value="$(arg role_name)"/>
+  </include>
+
+  <include file="$(find acting)/launch/acting.launch">
+      <arg name="control_loop_rate" value="$(arg control_loop_rate)"/>
+      <arg name="role_name" value="$(arg role_name)"/>
+  </include>
+  -->
+
+  <node pkg="carla_manual_control" type="carla_manual_control.py" name="carla_manual_control_$(arg role_name)" output="screen">
+    <param name="role_name" value="$(arg role_name)"/>
+  </node>
+
+
+  <node pkg="rqt_gui" type="rqt_gui" name="rqt_gui"/>
+  <node type="rviz" name="rviz" pkg="rviz" args="-d $(find agent)/config/rviz_config.rviz" />
+</launch>

code/perception/CMakeLists.txt

@@ -52,6 +52,7 @@ add_message_files(
   FILES
   Waypoint.msg
   LaneChange.msg
+  TrafficLightState.msg
   MinDistance.msg
 )
 

code/perception/launch/perception.launch

@@ -2,7 +2,7 @@
 <launch>
   <arg name="role_name" default="hero" />
   <arg name="control_loop_rate" default="0.1" />
-  
+
   <!--just for testing purposes -> Uncomment to use and see plots
   <node pkg="perception" type="sensor_filter_debug.py" name="sensor_filter_debug" output="screen">
     <param name="control_loop_rate" value="0.001" />
@@ -28,8 +28,7 @@
     <param name="control_loop_rate" value="0.1" />
     <param name="role_name" value="$(arg role_name)" />
   </node>
-
-
+
   <node pkg="perception" type="vision_node.py" name="VisionNode" output="screen">
     <param name="role_name" value="$(arg role_name)" />
     <param name="side" value="Center" />
@@ -52,9 +51,16 @@
 
       Image-Segmentation:
       - deeplabv3_resnet101
-      - yolov8x-seg -->
-
-    <param name="model" value="yolov8x-seg" />
+      - yolov8x-seg
+    -->
+
+    <param name="model" value="rtdetr-x" />
+  </node>
+
+  <node pkg="perception" type="traffic_light_node.py" name="TrafficLightNode" output="screen">
+    <param name="role_name" value="$(arg role_name)" />
+    <param name="side" value="Center" />
+    <param name="model" value="/workspace/code/perception/src/traffic_light_detection/models/model_acc_92.48_val_91.88.pt" />
   </node>
 
 

code/perception/msg/TrafficLightState.msg

@@ -0,0 +1,5 @@
+int8 state
+int8 GREEN=1
+int8 YELLOW=4
+int8 RED=2
+int8 UNKNOWN=0

code/perception/src/traffic_light_detection/.gitignore

@@ -1,2 +1 @@
 /dataset
-/models

code/perception/src/traffic_light_detection/dataset.dvc

@@ -1,5 +1,5 @@
 outs:
-- md5: 3a559397ebc58c1ecf142dea18d03367.dir
+- md5: 86f14bb96e1ac5735051b8e873f07a9f.dir
   size: 13745063
   nfiles: 2723
   hash: md5

code/perception/src/traffic_light_detection/dvc.lock

@@ -4,31 +4,32 @@ stages:
     cmd: python src/traffic_light_detection/traffic_light_training.py
     deps:
     - path: dataset
-      md5: 3a559397ebc58c1ecf142dea18d03367.dir
+      hash: md5
+      md5: 86f14bb96e1ac5735051b8e873f07a9f.dir
       size: 13745063
       nfiles: 2723
     - path: src
       hash: md5
-      md5: b6c9cb867c89ad6e86403d9c33538136.dir
-      size: 23777
-      nfiles: 10
+      md5: 34c981a61e886a858d135daee17a82ba.dir
+      size: 35849
+      nfiles: 17
     params:
       params.yaml:
         train:
-          epochs: 100
+          epochs: 20
           batch_size: 32
     outs:
     - path: dvclive/metrics.json
       hash: md5
-      md5: af33de699558fbfd3edee1607ba88f81
-      size: 218
+      md5: 8566265bcdc76cb55d17230f82fc1517
+      size: 219
     - path: dvclive/plots
       hash: md5
-      md5: 774919de9e9d6820ac6821d0819829c1.dir
-      size: 8900
+      md5: f8f99f42fc42e0ed3c80c8e8f05c1528.dir
+      size: 8870
       nfiles: 4
     - path: models
       hash: md5
-      md5: ee67bac2f189d2cc5a199d91ba3295ac.dir
-      size: 10815
+      md5: 16f96ecc475d20123051719908af9d4d.dir
+      size: 11071
       nfiles: 1

code/perception/src/traffic_light_detection/dvclive/metrics.json

@@ -0,0 +1,11 @@
+{
+    "train": {
+        "accuracy": 92.26495726495727,
+        "loss": 0.03128106306251298
+    },
+    "validation": {
+        "accuracy": 91.36125654450262,
+        "loss": 0.031184170882739323
+    },
+    "step": 99
+}

code/perception/src/traffic_light_detection/params.yaml

@@ -1,3 +1,3 @@
 train:
-  epochs: 100
+  epochs: 20
   batch_size: 32

code/perception/src/traffic_light_detection/src/traffic_light_config.py

@@ -17,7 +17,7 @@ def __init__(self):
         # Amount of epochs to train
         # One epoch: Training with all images from training dataset once
         self.NUM_WORKERS = 4
-        self.NUM_CLASSES = 4  # Traffic light states: green, yellow, red, back
+        self.NUM_CLASSES = 5  # States: green, yellow, red, back, side
         self.NUM_CHANNELS = 3  # RGB encoded images
 
         # Inference

code/perception/src/traffic_light_detection/src/traffic_light_detection/traffic_light_inference.py

@@ -2,11 +2,12 @@
 
 import torch.cuda
 import torchvision.transforms as t
-from data_generation.transforms import Normalize, ResizeAndPadToSquare, \
-    load_image
-from traffic_light_detection.classification_model import ClassificationModel
+from traffic_light_detection.src.traffic_light_detection.transforms \
+    import Normalize, ResizeAndPadToSquare, load_image
+from traffic_light_detection.src.traffic_light_detection.classification_model \
+    import ClassificationModel
 from torchvision.transforms import ToTensor
-from traffic_light_config import TrafficLightConfig
+from traffic_light_detection.src.traffic_light_config import TrafficLightConfig
 
 
 def parse_args():
@@ -49,7 +50,8 @@ def __init__(self, model_path):
         self.class_dict = {0: 'Backside',
                            1: 'Green',
                            2: 'Red',
-                           3: 'Yellow'}
+                           3: 'Side',
+                           4: 'Yellow'}
 
     def __call__(self, img):
         """

code/perception/src/traffic_light_detection/src/traffic_light_detection/traffic_light_training.py

@@ -11,8 +11,8 @@
 import sys
 import os
 sys.path.append(os.path.abspath(sys.path[0] + '/..'))
-from data_generation.transforms import Normalize, ResizeAndPadToSquare, \
-    load_image  # noqa: E402
+from traffic_light_detection.transforms import Normalize, \
+    ResizeAndPadToSquare, load_image  # noqa: E402
 from data_generation.weights_organizer import WeightsOrganizer  # noqa: E402
 from traffic_light_detection.classification_model import ClassificationModel \
     # noqa: E402

code/perception/src/traffic_light_node.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+
+from ros_compatibility.node import CompatibleNode
+import ros_compatibility as roscomp
+from rospy.numpy_msg import numpy_msg
+from sensor_msgs.msg import Image as ImageMsg
+from perception.msg import TrafficLightState
+from cv_bridge import CvBridge
+from traffic_light_detection.src.traffic_light_detection.traffic_light_inference \
+    import TrafficLightInference  # noqa: E501
+
+
+class TrafficLightNode(CompatibleNode):
+    def __init__(self, name, **kwargs):
+        super().__init__(name, **kwargs)
+        # general setup
+        self.bridge = CvBridge()
+        self.role_name = self.get_param("role_name", "hero")
+        self.side = self.get_param("side", "Center")
+        self.classifier = TrafficLightInference(self.get_param("model", ""))
+
+        # publish / subscribe setup
+        self.setup_camera_subscriptions()
+        self.setup_traffic_light_publishers()
+
+    def setup_camera_subscriptions(self):
+        self.new_subscription(
+            msg_type=numpy_msg(ImageMsg),
+            callback=self.handle_camera_image,
+            topic=f"/paf/{self.role_name}/{self.side}/segmented_traffic_light",
+            qos_profile=1
+        )
+
+    def setup_traffic_light_publishers(self):
+        self.traffic_light_publisher = self.new_publisher(
+            msg_type=TrafficLightState,
+            topic=f"/paf/{self.role_name}/{self.side}/traffic_light_state",
+            qos_profile=1
+        )
+
+    def handle_camera_image(self, image):
+        result = self.classifier(self.bridge.imgmsg_to_cv2(image))
+
+        # 1: Green, 2: Red, 4: Yellow, 0: Unknown
+        msg = TrafficLightState()
+        msg.state = result if result in [1, 2, 4] else 0
+
+        self.traffic_light_publisher.publish(msg)
+
+    def run(self):
+        self.spin()
+
+
+if __name__ == "__main__":
+    roscomp.init("TrafficLightNode")
+    node = TrafficLightNode("TrafficLightNode")
+    node.run()

code/perception/src/vision_node.py

@@ -85,6 +85,7 @@ def __init__(self, name, **kwargs):
         # publish / subscribe setup
         self.setup_camera_subscriptions()
         self.setup_camera_publishers()
+        self.setup_traffic_light_publishers()
         self.image_msg_header = Header()
         self.image_msg_header.frame_id = "segmented_image_frame"
 
@@ -127,6 +128,13 @@ def setup_camera_publishers(self):
             qos_profile=1
         )
 
+    def setup_traffic_light_publishers(self):
+        self.traffic_light_publisher = self.new_publisher(
+            msg_type=numpy_msg(ImageMsg),
+            topic=f"/paf/{self.role_name}/{self.side}/segmented_traffic_light",
+            qos_profile=1
+        )
+
     def handle_camera_image(self, image):
         # free up cuda memory
         if self.device == "cuda":
@@ -140,12 +148,10 @@ def handle_camera_image(self, image):
 
         # publish image to rviz
         img_msg = self.bridge.cv2_to_imgmsg(vision_result,
-                                            encoding="passthrough")
+                                            encoding="rgb8")
         img_msg.header = image.header
         self.publisher.publish(img_msg)
 
-        pass
-
     def predict_torch(self, image):
         self.model.eval()
         cv_image = self.bridge.imgmsg_to_cv2(img_msg=image,
@@ -174,10 +180,35 @@ def predict_ultralytics(self, image):
                                              desired_encoding='passthrough')
         cv_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2BGR)
 
-        output = self.model(cv_image)
+        output = self.model(cv_image, half=True, verbose=False)
+
+        if 9 in output[0].boxes.cls:
+            self.process_traffic_lights(output[0], cv_image, image.header)
 
         return output[0].plot()
 
+    def process_traffic_lights(self, prediction, cv_image, image_header):
+        indices = (prediction.boxes.cls == 9).nonzero().squeeze().cpu().numpy()
+        indices = np.asarray([indices]) if indices.size == 1 else indices
+
+        min_x = 550
+        max_x = 700
+        min_prob = 0.35
+
+        for index in indices:
+            box = prediction.boxes.cpu().data.numpy()[index]
+
+            if box[0] < min_x or box[2] > max_x or box[4] < min_prob:
+                continue
+
+            box = box[0:4].astype(int)
+            segmented = cv_image[box[1]:box[3], box[0]:box[2]]
+
+            traffic_light_image = self.bridge.cv2_to_imgmsg(segmented,
+                                                            encoding="rgb8")
+            traffic_light_image.header = image_header
+            self.traffic_light_publisher.publish(traffic_light_image)
+
     def create_mask(self, input_image, model_output):
         output_predictions = torch.argmax(model_output, dim=0)
         for i in range(21):

code/planning/src/behavior_agent/behavior_tree.py

@@ -1,21 +1,19 @@
 #!/usr/bin/env python
 
 import functools
-# import behavior_agent
-import py_trees
 from py_trees.behaviours import Running
 import py_trees_ros
-import py_trees.console as console
 import rospy
 import sys
 import behaviours
 from py_trees.composites import Parallel, Selector, Sequence
-from py_trees.decorators import Inverter
 
 """
 Source: https://github.com/ll7/psaf2
 """
 
+# flake8: noqa: E501
+
 
 def grow_a_tree(role_name):
 
@@ -61,7 +59,7 @@ def grow_a_tree(role_name):
                                                  ("Leave Change")
                                              ])
                                 ]),
-                        
+
                     ]),
                     behaviours.maneuvers.Cruise("Cruise")
                 ])

code/planning/src/behavior_agent/behaviours/__init__.py

@@ -1,2 +0,0 @@
-from . import topics2blackboard, road_features
-from . import intersection, traffic_objects, maneuvers, meta, lane_change