fix(autoware_traffic_light_visualization): fix to visualize correct color and shapes

perception/autoware_traffic_light_visualization/src/traffic_light_roi_visualizer/node.cpp

@@ -1,4 +1,4 @@
 // Copyright 2020 Tier IV, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -20,7 +20,6 @@
 
 #include <memory>
 #include <string>
-#include <utility>
 
 namespace autoware::traffic_light
 {
@@ -103,26 +102,18 @@
   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{254, 149, 149};
-  } else if (result.label.find("yellow") != std::string::npos) {
-    color = cv::Scalar{254, 250, 149};
-  } else if (result.label.find("green") != std::string::npos) {
-    color = cv::Scalar{149, 254, 161};
-  } else {
-    color = cv::Scalar{250, 250, 250};
-  }
+  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, 2);
+    info.color, 2);
 
-  std::string shape_name = extractShapeName(result.label);
+  constexpr int shape_img_size = 16;
+  const auto position = cv::Point(tl_roi.roi.x_offset, tl_roi.roi.y_offset);
 
   visualization::drawTrafficLightShape(
-    image, shape_name, cv::Point(tl_roi.roi.x_offset, tl_roi.roi.y_offset), color, 16, result.prob);
+    image, info.shapes, shape_img_size, position, info.color, result.prob);
 
   return true;
 }

perception/autoware_traffic_light_visualization/src/traffic_light_roi_visualizer/node.hpp

@@ -36,7 +36,9 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <sstream>
 #include <string>
+#include <vector>
 
 namespace autoware::traffic_light
 {
@@ -46,6 +48,15 @@
   std::string label;
 };
 
+/**
+ * @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:
@@ -78,25 +89,57 @@
     {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"},
   };
 
-  std::string extractShapeName(const std::string & label)
+  /**
+   * @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)
   {
-    size_t start_pos = label.find('-');
-    if (start_pos != std::string::npos) {
-      start_pos++;                                  // Start after the hyphen
-      size_t end_pos = label.find(',', start_pos);  // Find the next comma after the hyphen
-      if (end_pos == std::string::npos) {  // If no comma is found, take the rest of the string
-        end_pos = label.length();
+    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 label.substr(start_pos, end_pos - start_pos);
     }
-    return "unknown";  // Return "unknown" if no hyphen is found
+    return {color, shapes};
   }
 
   bool createRect(

perception/autoware_traffic_light_visualization/src/traffic_light_roi_visualizer/shape_draw.cpp

@@ -17,164 +17,128 @@
 #include "opencv2/highgui.hpp"
 #include "opencv2/imgproc.hpp"
 
+#include <algorithm>
+#include <cmath>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
 namespace autoware::traffic_light::visualization
 {
 void drawShape(
-  const DrawFunctionParams & params, const std::string & filename, bool flipHorizontally,
-  bool flipVertically, int x_offset, int y_offset, double scale_factor)
+  cv::Mat & image, const std::vector<ShapeImgParam> & params, int size, const cv::Point & position,
+  const cv::Scalar & color, float probability)
 {
-  std::string filepath =
-    ament_index_cpp::get_package_share_directory("autoware_traffic_light_visualization") +
-    "/images/" + filename;
-  cv::Mat shapeImg = cv::imread(filepath, cv::IMREAD_UNCHANGED);
-  if (shapeImg.empty()) {
-    std::cerr << "Failed to load image: " << filepath << std::endl;
-    return;
-  }
+  // load concatenated shape image
+  const auto shape_img = loadShapeImage(params, size);
 
-  if (flipHorizontally) {
-    cv::flip(shapeImg, shapeImg, 1);  // Flip horizontally
-  }
+  // Calculate the width of the text
+  std::string prob_str = std::to_string(static_cast<int>(round(probability * 100))) + "%";
 
-  if (flipVertically) {
-    cv::flip(shapeImg, shapeImg, 0);  // Flip vertically
-  }
+  int baseline = 0;
+  cv::Size text_size = cv::getTextSize(prob_str, cv::FONT_HERSHEY_SIMPLEX, 0.75, 2, &baseline);
 
-  cv::resize(
-    shapeImg, shapeImg, cv::Size(params.size, params.size), scale_factor, scale_factor,
-    cv::INTER_AREA);
+  const int fill_rect_w = shape_img.cols + text_size.width + 20;
+  const int fill_rect_h = std::max(shape_img.rows, text_size.height) + 10;
 
-  // Calculate the center position including offsets
-  cv::Point position(
-    params.position.x + x_offset, params.position.y - shapeImg.rows / 2 + y_offset);
+  const cv::Point rect_position(position.x, position.y - fill_rect_h);
 
-  // Check for image boundaries
   if (
-    position.x < 0 || position.y < 0 || position.x + shapeImg.cols > params.image.cols ||
-    position.y + shapeImg.rows > params.image.rows) {
+    rect_position.x < 0 || rect_position.y < 0 || rect_position.x + fill_rect_w > image.cols ||
+    position.y > image.rows) {
     // TODO(KhalilSelyan): This error message may flood the terminal logs, so commented out
     // temporarily. Need to consider a better way.
 
     // std::cerr << "Adjusted position is out of image bounds." << std::endl;
     return;
   }
-
-  // Calculate the width of the text
-  std::string probabilityText =
-    std::to_string(static_cast<int>(round(params.probability * 100))) + "%";
-  int baseline = 0;
-  cv::Size textSize =
-    cv::getTextSize(probabilityText, cv::FONT_HERSHEY_SIMPLEX, 0.75, 2, &baseline);
-
-  // Adjust the filled rectangle to be at the top edge and the correct width
-  int filledRectWidth =
-    shapeImg.cols + (filename != "unknown.png" ? textSize.width + 10 : 5);  // Add some padding
-  int filledRectHeight = shapeImg.rows + 10;                                // Add some padding
-
-  cv::rectangle(
-    params.image, cv::Rect(position.x - 2, position.y - 5, filledRectWidth, filledRectHeight),
-    params.color,
-    -1);  // Filled rectangle
-
-  // Create ROI on the destination image
-  cv::Mat destinationROI = params.image(cv::Rect(position, cv::Size(shapeImg.cols, shapeImg.rows)));
-
-  // Overlay the image onto the main image
-  for (int y = 0; y < shapeImg.rows; ++y) {
-    for (int x = 0; x < shapeImg.cols; ++x) {
-      cv::Vec4b & pixel = shapeImg.at<cv::Vec4b>(y, x);
-      if (pixel[3] != 0) {  // Only non-transparent pixels
-        destinationROI.at<cv::Vec3b>(y, x) = cv::Vec3b(pixel[0], pixel[1], pixel[2]);
+  cv::Rect rectangle(rect_position.x, rect_position.y, fill_rect_w, fill_rect_h);
+  cv::rectangle(image, rectangle, color, -1);
+
+  // Position the probability text right next to the shape. (Text origin: bottom-left)
+  const int prob_x_offset = shape_img.cols + 15;
+  const int prob_y_offset = fill_rect_h / 4;
+  cv::putText(
+    image, prob_str, cv::Point(position.x + prob_x_offset, position.y - prob_y_offset),
+    cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(0, 0, 0), 2, cv::LINE_AA);
+
+  if (!shape_img.empty()) {
+    // Create ROI on the destination image
+    const int shape_y_offset = fill_rect_h / 4;
+    auto shapeRoi = image(cv::Rect(
+      rect_position.x + 5, rect_position.y + shape_y_offset, shape_img.cols, shape_img.rows));
+
+    // Overlay the image onto the main image
+    for (int y = 0; y < shape_img.rows; ++y) {
+      for (int x = 0; x < shape_img.cols; ++x) {
+        const auto & pixel = shape_img.at<cv::Vec4b>(y, x);
+        if (pixel[3] != 0) {  // Only non-transparent pixels
+          shapeRoi.at<cv::Vec3b>(y, x) = cv::Vec3b(pixel[0], pixel[1], pixel[2]);
+        }
       }
     }
   }
-
-  // Position the probability text right next to the shape
-  if (filename != "unknown.png") {
-    cv::putText(
-      params.image, probabilityText,
-      cv::Point(
-        position.x + shapeImg.cols + 5, position.y + shapeImg.rows / 2 + textSize.height / 2),
-      cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(0, 0, 0), 2, cv::LINE_AA);
-  }
-}
-
-void drawCircle(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
-  drawShape(params, "circle.png", false, false, 0, -y_offset);
-}
-
-void drawLeftArrow(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
-  drawShape(params, "left_arrow.png", false, false, 0, -y_offset);
-}
-
-void drawRightArrow(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
-  drawShape(params, "left_arrow.png", true, false, 0, -y_offset);
 }
 
-void drawStraightArrow(const DrawFunctionParams & params)
+cv::Mat loadShapeImage(const std::vector<ShapeImgParam> & params, int size, double scale_factor)
 {
-  int y_offset = params.size / 2 + 5;  // This adjusts the base position upwards
+  if (params.empty()) {
+    return {};
+  }
 
-  drawShape(params, "straight_arrow.png", false, false, 0, -y_offset);
-}
-void drawDownArrow(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;  // This adjusts the base position upwards
-  drawShape(params, "straight_arrow.png", false, true, 0, -y_offset);
-}
+  static const auto img_dir =
+    ament_index_cpp::get_package_share_directory("autoware_traffic_light_visualization") +
+    "/images/";
 
-void drawDownLeftArrow(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
-  drawShape(params, "down_left_arrow.png", false, false, 0, -y_offset);
-}
+  std::vector<cv::Mat> src_img;
+  for (const auto & param : params) {
+    auto filepath = img_dir + param.filename;
+    auto img = cv::imread(filepath, cv::IMREAD_UNCHANGED);
 
-void drawDownRightArrow(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
-  drawShape(params, "down_left_arrow.png", true, false, 0, -y_offset);
-}
+    cv::resize(img, img, cv::Size(size, size), scale_factor, scale_factor, cv::INTER_AREA);
 
-void drawCross(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
+    if (param.h_flip) {
+      cv::flip(img, img, 1);
+    }
+    if (param.v_flip) {
+      cv::flip(img, img, 0);
+    }
+    src_img.emplace_back(img);
+  }
 
-  drawShape(params, "cross.png", false, false, 0, -y_offset);
-}
+  cv::Mat dst;
+  cv::hconcat(src_img, dst);  // cspell:ignore hconcat
 
-void drawUnknown(const DrawFunctionParams & params)
-{
-  int y_offset = params.size / 2 + 5;
-  drawShape(params, "unknown.png", false, false, 0, -y_offset);
+  return dst;
 }
 
 void drawTrafficLightShape(
-  cv::Mat & image, const std::string & shape, const cv::Point & position, const cv::Scalar & color,
-  int size, float probability)
+  cv::Mat & image, const std::vector<std::string> & shapes, int size, const cv::Point & position,
+  const cv::Scalar & color, float probability)
 {
-  static std::map<std::string, DrawFunction> shapeToFunction = {
-    {"circle", drawCircle},
-    {"left", drawLeftArrow},
-    {"right", drawRightArrow},
-    {"straight", drawStraightArrow},
-    {"down", drawDownArrow},
-    {"down_left", drawDownLeftArrow},
-    {"down_right", drawDownRightArrow},
-    {"cross", drawCross},
-    {"unknown", drawUnknown}};
-  auto it = shapeToFunction.find(shape);
-  if (it != shapeToFunction.end()) {
-    DrawFunctionParams params{image, position, color, size, probability};
-    it->second(params);
-  } else {
-    std::cerr << "Unknown shape: " << shape << std::endl;
+  using ShapeImgParamFunction = std::function<ShapeImgParam()>;
+
+  static const std::unordered_map<std::string, ShapeImgParamFunction> shapeToParamFunction = {
+    {"circle", circleImgParam},
+    {"left", leftArrowImgParam},
+    {"right", rightArrowImgParam},
+    {"straight", straightArrowImgParam},
+    {"down", downArrowImgParam},
+    {"straight_left", straightLeftArrowImgParam},
+    {"straight_right", straightRightArrowImgParam},
+    {"down_left", downLeftArrowImgParam},
+    {"down_right", downRightArrowImgParam},
+    {"cross", crossImgParam},
+    {"unknown", unknownImgParam}};
+
+  std::vector<ShapeImgParam> params;
+  for (const auto & shape : shapes) {
+    if (shapeToParamFunction.find(shape) != shapeToParamFunction.end()) {
+      auto func = shapeToParamFunction.at(shape);
+      params.emplace_back(func());
+    }
   }
-}
 
+  drawShape(image, params, size, position, color, probability);
+}
 }  // namespace autoware::traffic_light::visualization