fix(autoware_traffic_light_visualization): fix to visualize correct color and shapes (autowarefoundation#8428)

fix(autoware_traffic_light_visualization): fix vialization to draw correct shapes

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>
Co-authored-by: Yi-Hsiang Fang (Vivid) <146902905+vividf@users.noreply.github.com>

Author: 2 people

perception/traffic_light_visualization/include/traffic_light_visualization/traffic_light_roi_visualizer/nodelet.hpp

@@ -32,7 +32,9 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <sstream>
 #include <string>
+#include <vector>
 
 namespace traffic_light
 {
@@ -42,7 +44,16 @@ struct ClassificationResult
   std::string label;
 };
 
-class TrafficLightRoiVisualizerNodelet : public rclcpp::Node
+/**
+ * @brief A struct to represent parsed traffic light shape information.
+ */
+struct TrafficLightShapeInfo
+{
+  cv::Scalar color;                 //!< Color associated with "circle".
+  std::vector<std::string> shapes;  //!< Shape names.
+};
+
+class TrafficLightRoiVisualizerNode : public rclcpp::Node
 {
 public:
   explicit TrafficLightRoiVisualizerNodelet(const rclcpp::NodeOptions & options);
@@ -74,13 +85,59 @@ class TrafficLightRoiVisualizerNodelet : public rclcpp::Node
     {tier4_perception_msgs::msg::TrafficLightElement::RIGHT_ARROW, "right"},
     {tier4_perception_msgs::msg::TrafficLightElement::UP_ARROW, "straight"},
     {tier4_perception_msgs::msg::TrafficLightElement::DOWN_ARROW, "down"},
+    {tier4_perception_msgs::msg::TrafficLightElement::UP_LEFT_ARROW, "straight_left"},
+    {tier4_perception_msgs::msg::TrafficLightElement::UP_RIGHT_ARROW, "straight_right"},
     {tier4_perception_msgs::msg::TrafficLightElement::DOWN_LEFT_ARROW, "down_left"},
     {tier4_perception_msgs::msg::TrafficLightElement::DOWN_RIGHT_ARROW, "down_right"},
     {tier4_perception_msgs::msg::TrafficLightElement::CROSS, "cross"},
     // other
     {tier4_perception_msgs::msg::TrafficLightElement::UNKNOWN, "unknown"},
   };
 
+  /**
+   * @brief Return RGB color from color string associated with "circle".
+   * @param color Color string.
+   * @return RGB color.
+   */
+  static cv::Scalar strToColor(const std::string & color)
+  {
+    if (color == "red") {
+      return {254, 149, 149};
+    } else if (color == "yellow") {
+      return {254, 250, 149};
+    } else if (color == "green") {
+      return {149, 254, 161};
+    } else {
+      return {250, 250, 250};
+    }
+  }
+
+  /**
+   * @brief Extract color and shape names from label.
+   * @param label String formatted as `<Color0>-<Shape0>,<Color1>-<Shape1>,...,<ColorN>-<ShapeN>`.
+   * @return Extracted information includes a color associated with "circle" and shape names.
+   */
+  static TrafficLightShapeInfo extractShapeInfo(const std::string & label)
+  {
+    cv::Scalar color{255, 255, 255};
+    std::vector<std::string> shapes;
+
+    std::stringstream ss(label);
+    std::string segment;
+    while (std::getline(ss, segment, ',')) {
+      size_t hyphen_pos = segment.find('-');
+      if (hyphen_pos != std::string::npos) {
+        auto shape = segment.substr(hyphen_pos + 1);
+        if (shape == "circle") {
+          const auto color_str = segment.substr(0, hyphen_pos);
+          color = strToColor(color_str);
+        }
+        shapes.emplace_back(shape);
+      }
+    }
+    return {color, shapes};
+  }
+
   bool createRect(
     cv::Mat & image, const tier4_perception_msgs::msg::TrafficLightRoi & tl_roi,
     const cv::Scalar & color);

perception/traffic_light_visualization/src/traffic_light_roi_visualizer/nodelet.cpp

@@ -19,7 +19,6 @@
 
 #include <memory>
 #include <string>
-#include <utility>
 
 namespace traffic_light
 {
@@ -91,31 +90,18 @@ bool TrafficLightRoiVisualizerNodelet::createRect(
   cv::Mat & image, const tier4_perception_msgs::msg::TrafficLightRoi & tl_roi,
   const ClassificationResult & result)
 {
-  cv::Scalar color;
-  if (result.label.find("red") != std::string::npos) {
-    color = cv::Scalar{255, 0, 0};
-  } else if (result.label.find("yellow") != std::string::npos) {
-    color = cv::Scalar{0, 255, 0};
-  } else if (result.label.find("green") != std::string::npos) {
-    color = cv::Scalar{0, 0, 255};
-  } else {
-    color = cv::Scalar{255, 255, 255};
-  }
+  const auto info = extractShapeInfo(result.label);
 
   cv::rectangle(
     image, cv::Point(tl_roi.roi.x_offset, tl_roi.roi.y_offset),
     cv::Point(tl_roi.roi.x_offset + tl_roi.roi.width, tl_roi.roi.y_offset + tl_roi.roi.height),
-    color, 3);
+    info.color, 2);
 
-  int offset = 40;
-  cv::putText(
-    image, std::to_string(result.prob),
-    cv::Point(tl_roi.roi.x_offset, tl_roi.roi.y_offset - (offset * 0)), cv::FONT_HERSHEY_COMPLEX,
-    1.1, color, 3);
+  constexpr int shape_img_size = 16;
+  const auto position = cv::Point(tl_roi.roi.x_offset, tl_roi.roi.y_offset);
 
-  cv::putText(
-    image, result.label, cv::Point(tl_roi.roi.x_offset, tl_roi.roi.y_offset - (offset * 1)),
-    cv::FONT_HERSHEY_COMPLEX, 1.1, color, 2);
+  visualization::drawTrafficLightShape(
+    image, info.shapes, shape_img_size, position, info.color, result.prob);
 
   return true;
 }