Fix segmentation type case and format

Signed-off-by: Florian Vahl <florian@flova.de>

Author: Flova

src/bitbots_motion/bitbots_head_mover/src/move_head.cpp

@@ -146,8 +146,7 @@ class HeadMover {
     // Bring the goal point into the planning frame
     geometry_msgs::msg::PointStamped new_point;
     try {
-      new_point =
-          tf_buffer_->transform(goal->look_at_position, "base_footprint", tf2::durationFromSec(0.9));
+      new_point = tf_buffer_->transform(goal->look_at_position, "base_footprint", tf2::durationFromSec(0.9));
     } catch (tf2::TransformException& ex) {
       RCLCPP_ERROR(node_->get_logger(), "Could not transform goal point: %s", ex.what());
       return rclcpp_action::GoalResponse::REJECT;

src/bitbots_vision/include/bitbots_vision/candidate.hpp

@@ -13,8 +13,7 @@ struct Candidate {
   int height{0};  ///< Box height
   float rating{0.f};
 
-  static Candidate from_x1y1x2y2(int x1, int y1, int x2, int y2, float rating)
-  {
+  static Candidate from_x1y1x2y2(int x1, int y1, int x2, int y2, float rating) {
     return {std::min(x1, x2), std::min(y1, y2), std::abs(x2 - x1), std::abs(y2 - y1), rating};
   }
 
@@ -24,18 +23,15 @@ struct Candidate {
   int y2() const { return y1 + height; }
   int radius() const { return (width + height) / 4; }
 
-  static std::vector<Candidate> sort_by_rating(std::vector<Candidate> candidates)
-  {
+  static std::vector<Candidate> sort_by_rating(std::vector<Candidate> candidates) {
     std::sort(candidates.begin(), candidates.end(),
-              [](const Candidate & a, const Candidate & b) { return a.rating > b.rating; });
+              [](const Candidate& a, const Candidate& b) { return a.rating > b.rating; });
     return candidates;
   }
 
-  static std::vector<Candidate> filter_by_rating(
-    const std::vector<Candidate> & candidates, float threshold)
-  {
+  static std::vector<Candidate> filter_by_rating(const std::vector<Candidate>& candidates, float threshold) {
     std::vector<Candidate> result;
-    for (const auto & c : candidates) {
+    for (const auto& c : candidates) {
       if (c.rating > threshold) {
         result.push_back(c);
       }

src/bitbots_vision/include/bitbots_vision/debug_image.hpp

@@ -7,40 +7,33 @@
 namespace bitbots_vision {
 
 /// Draws detection and segmentation results on a copy of the input image.
-class DebugImage
-{
-public:
+class DebugImage {
+ public:
   // BGR color constants (cv::Scalar is not a literal type, so inline static const)
-  inline static const cv::Scalar kBall{0, 255, 0};             // green
+  inline static const cv::Scalar kBall{0, 255, 0};                // green
   inline static const cv::Scalar kRobotTeamMates{255, 255, 102};  // cyan
   inline static const cv::Scalar kRobotOpponents{153, 51, 255};   // magenta
   inline static const cv::Scalar kRobotUnknown{160, 160, 160};    // grey
-  inline static const cv::Scalar kGoalposts{255, 255, 255};        // white
-  inline static const cv::Scalar kLines{255, 0, 0};                // blue
+  inline static const cv::Scalar kGoalposts{255, 255, 255};       // white
+  inline static const cv::Scalar kLines{255, 0, 0};               // blue
 
   explicit DebugImage(bool active = false) : active_(active) {}
 
-  void set_image(const cv::Mat & image);
+  void set_image(const cv::Mat& image);
 
   /// Draw a circle around each candidate (suitable for ball-shaped objects).
-  void draw_ball_candidates(
-    const std::vector<Candidate> & candidates,
-    const cv::Scalar & color,
-    int thickness = 1);
+  void draw_ball_candidates(const std::vector<Candidate>& candidates, const cv::Scalar& color, int thickness = 1);
 
   /// Draw a bounding rectangle for each candidate (suitable for robots, goalposts).
-  void draw_box_candidates(
-    const std::vector<Candidate> & candidates,
-    const cv::Scalar & color,
-    int thickness = 1);
+  void draw_box_candidates(const std::vector<Candidate>& candidates, const cv::Scalar& color, int thickness = 1);
 
   /// Blend a segmentation mask over the debug image.
   /// @param mask  CV_8UC1 binary mask (non-zero = active)
-  void draw_mask(const cv::Mat & mask, const cv::Scalar & color, double opacity = 0.5);
+  void draw_mask(const cv::Mat& mask, const cv::Scalar& color, double opacity = 0.5);
 
-  const cv::Mat & get_image() const { return image_; }
+  const cv::Mat& get_image() const { return image_; }
 
-private:
+ private:
   bool active_{false};
   cv::Mat image_;
 };

src/bitbots_vision/include/bitbots_vision/model_config.hpp

@@ -6,22 +6,21 @@
 namespace bitbots_vision {
 
 /// Configuration loaded from the model's `model_config.yaml`.
-class ModelConfig
-{
-public:
+class ModelConfig {
+ public:
   ModelConfig() = default;
 
   /// Load from `<model_dir>/model_config.yaml`.  Throws on error.
-  static ModelConfig load_from(const std::string & model_dir);
+  static ModelConfig load_from(const std::string& model_dir);
 
-  const std::vector<std::string> & detection_classes() const { return detection_classes_; }
-  const std::vector<std::string> & segmentation_classes() const { return segmentation_classes_; }
+  const std::vector<std::string>& detection_classes() const { return detection_classes_; }
+  const std::vector<std::string>& segmentation_classes() const { return segmentation_classes_; }
   bool team_colors_provided() const { return team_colors_provided_; }
 
   /// Indices of all classes whose name contains "robot".
   std::vector<int> robot_class_ids() const;
 
-private:
+ private:
   std::vector<std::string> detection_classes_;
   std::vector<std::string> segmentation_classes_;
   bool team_colors_provided_{false};

src/bitbots_vision/include/bitbots_vision/yoeo_handler.hpp

@@ -1,14 +1,14 @@
 #pragma once
 
+#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
+
 #include <memory>
+#include <opencv2/core.hpp>
+#include <rclcpp/logger.hpp>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
-#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
-#include <opencv2/core.hpp>
-#include <rclcpp/logger.hpp>
-
 #include "bitbots_vision/candidate.hpp"
 #include "bitbots_vision/model_config.hpp"
 #include "bitbots_vision/yoeo_processing.hpp"
@@ -25,44 +25,40 @@ namespace bitbots_vision {
 ///   - two outputs:
 ///       [0] detections     shape [1, N, 5+num_classes]  (x_c, y_c, w, h, obj_conf, class_probs…)
 ///       [1] segmentation   shape [1, H_seg, W_seg]      (argmax class index per pixel)
-class YoeoHandler
-{
-public:
+class YoeoHandler {
+ public:
   struct Config {
     float conf_threshold{0.5f};
     float nms_threshold{0.4f};
   };
 
-  YoeoHandler(
-    const std::string & model_path,
-    const ModelConfig & model_config,
-    const Config & cfg,
-    const rclcpp::Logger & logger);
+  YoeoHandler(const std::string& model_path, const ModelConfig& model_config, const Config& cfg,
+              const rclcpp::Logger& logger);
 
   /// Update thresholds without reloading the model.
-  void reconfigure(const Config & cfg);
+  void reconfigure(const Config& cfg);
 
   /// Set the image to be processed by the next call to predict().
   /// The image must be in BGR8 format (as delivered by cv_bridge).
-  void set_image(const cv::Mat & bgr_image);
+  void set_image(const cv::Mat& bgr_image);
 
   /// Run the network on the current image (no-op if already up to date).
   void predict();
 
-  std::vector<Candidate> get_detection_candidates_for(const std::string & class_name);
-  cv::Mat get_segmentation_mask_for(const std::string & class_name);
+  std::vector<Candidate> get_detection_candidates_for(const std::string& class_name);
+  cv::Mat get_segmentation_mask_for(const std::string& class_name);
 
-  const std::vector<std::string> & detection_class_names() const;
-  const std::vector<std::string> & segmentation_class_names() const;
+  const std::vector<std::string>& detection_class_names() const;
+  const std::vector<std::string>& segmentation_class_names() const;
 
-private:
+ private:
   // ----- ONNX Runtime objects -----
   Ort::Env env_;
   Ort::SessionOptions session_options_;
   std::unique_ptr<Ort::Session> session_;
 
   // ----- Model metadata -----
-  std::vector<int64_t> input_shape_;   // [1, 3, H_net, W_net]
+  std::vector<int64_t> input_shape_;  // [1, 3, H_net, W_net]
   std::string input_name_;
   std::vector<std::string> output_names_;
   int num_det_classes_{0};
@@ -74,7 +70,7 @@ class YoeoHandler
   Config cfg_;
   rclcpp::Logger logger_;
 
-  cv::Mat current_image_rgb_;      // float32, CHW layout, stored as [C, H, W] in flat vector
+  cv::Mat current_image_rgb_;  // float32, CHW layout, stored as [C, H, W] in flat vector
   std::vector<float> input_data_;
 
   bool prediction_is_fresh_{true};
@@ -84,14 +80,14 @@ class YoeoHandler
 
   // Cached results
   std::unordered_map<std::string, std::vector<Candidate>> det_results_;
-  std::unordered_map<std::string, cv::Mat> seg_results_;   // CV_8UC1 binary masks
+  std::unordered_map<std::string, cv::Mat> seg_results_;  // CV_8UC1 binary masks
 
   // ----- Helpers -----
-  void init_session(const std::string & model_path);
-  void preprocess(const cv::Mat & bgr_image);
+  void init_session(const std::string& model_path);
+  void preprocess(const cv::Mat& bgr_image);
   void run_inference();
-  void postprocess_detections(const float * det_data, const std::vector<int64_t> & shape);
-  void postprocess_segmentation(const float * seg_data, const std::vector<int64_t> & shape);
+  void postprocess_detections(const float* det_data, const std::vector<int64_t>& shape);
+  void postprocess_segmentation(const uint8_t* seg_data, const std::vector<int64_t>& shape);
 };
 
 }  // namespace bitbots_vision

src/bitbots_vision/include/bitbots_vision/yoeo_processing.hpp

@@ -1,11 +1,10 @@
 #pragma once
 
+#include <opencv2/core.hpp>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
-#include <opencv2/core.hpp>
-
 #include "bitbots_vision/candidate.hpp"
 
 /// Pure preprocessing / postprocessing functions for YOEO inference.
@@ -23,7 +22,7 @@ namespace bitbots_vision::processing {
 struct PreprocessInfo {
   int orig_h{0};
   int orig_w{0};
-  int max_dim{0};      ///< max(orig_h, orig_w) == side length of the padded square
+  int max_dim{0};  ///< max(orig_h, orig_w) == side length of the padded square
   int pad_top{0};
   int pad_bottom{0};
   int pad_left{0};
@@ -46,8 +45,7 @@ struct PreprocessInfo {
 /// @param net_w    Network input width
 /// @param info     [out] Padding / sizing metadata for use in postprocessors
 /// @return         Flattened CHW float32 tensor of size 3 * net_h * net_w
-std::vector<float> preprocess_image(
-  const cv::Mat & bgr, int net_h, int net_w, PreprocessInfo & info);
+std::vector<float> preprocess_image(const cv::Mat& bgr, int net_h, int net_w, PreprocessInfo& info);
 
 // ---------------------------------------------------------------------------
 // Non-maximum suppression
@@ -66,13 +64,9 @@ std::vector<float> preprocess_image(
 /// @param nms_threshold    IoU threshold for suppression
 /// @param max_detections   Upper limit on kept boxes
 /// @return                 Indices of kept boxes
-std::vector<int> nms_boxes(
-  const std::vector<cv::Rect2d> & boxes,
-  const std::vector<float> & scores,
-  const std::vector<int> & class_ids,
-  const std::vector<int> & robot_class_ids,
-  float nms_threshold,
-  int max_detections = 30);
+std::vector<int> nms_boxes(const std::vector<cv::Rect2d>& boxes, const std::vector<float>& scores,
+                           const std::vector<int>& class_ids, const std::vector<int>& robot_class_ids,
+                           float nms_threshold, int max_detections = 30);
 
 // ---------------------------------------------------------------------------
 // Detection postprocessing
@@ -94,14 +88,8 @@ std::vector<int> nms_boxes(
 /// @param info             Preprocessing metadata for coordinate rescaling
 /// @return                 Map of class_name → kept Candidates
 std::unordered_map<std::string, std::vector<Candidate>> postprocess_detections(
-  const float * data,
-  int64_t num_boxes,
-  int64_t stride,
-  const std::vector<std::string> & class_names,
-  const std::vector<int> & robot_class_ids,
-  float conf_thresh,
-  float nms_thresh,
-  const PreprocessInfo & info);
+    const float* data, int64_t num_boxes, int64_t stride, const std::vector<std::string>& class_names,
+    const std::vector<int>& robot_class_ids, float conf_thresh, float nms_thresh, const PreprocessInfo& info);
 
 // ---------------------------------------------------------------------------
 // Segmentation postprocessing
@@ -122,11 +110,8 @@ std::unordered_map<std::string, std::vector<Candidate>> postprocess_detections(
 /// @param class_names  Segmentation class names (index == class id)
 /// @param info         Preprocessing metadata for unpadding
 /// @return             Map of class_name → CV_8UC1 binary mask
-std::unordered_map<std::string, cv::Mat> postprocess_segmentation(
-  const float * data,
-  int seg_h,
-  int seg_w,
-  const std::vector<std::string> & class_names,
-  const PreprocessInfo & info);
+std::unordered_map<std::string, cv::Mat> postprocess_segmentation(const uint8_t* data, int seg_h, int seg_w,
+                                                                  const std::vector<std::string>& class_names,
+                                                                  const PreprocessInfo& info);
 
 }  // namespace bitbots_vision::processing

src/bitbots_vision/src/debug_image.cpp

@@ -4,39 +4,28 @@
 
 namespace bitbots_vision {
 
-void DebugImage::set_image(const cv::Mat & image)
-{
-  image_ = image.clone();
-}
+void DebugImage::set_image(const cv::Mat& image) { image_ = image.clone(); }
 
-void DebugImage::draw_ball_candidates(
-  const std::vector<Candidate> & candidates,
-  const cv::Scalar & color,
-  int thickness)
-{
+void DebugImage::draw_ball_candidates(const std::vector<Candidate>& candidates, const cv::Scalar& color,
+                                      int thickness) {
   if (!active_) {
     return;
   }
-  for (const auto & c : candidates) {
+  for (const auto& c : candidates) {
     cv::circle(image_, {c.center_x(), c.center_y()}, c.radius(), color, thickness);
   }
 }
 
-void DebugImage::draw_box_candidates(
-  const std::vector<Candidate> & candidates,
-  const cv::Scalar & color,
-  int thickness)
-{
+void DebugImage::draw_box_candidates(const std::vector<Candidate>& candidates, const cv::Scalar& color, int thickness) {
   if (!active_) {
     return;
   }
-  for (const auto & c : candidates) {
+  for (const auto& c : candidates) {
     cv::rectangle(image_, {c.x1, c.y1}, {c.x2(), c.y2()}, color, thickness);
   }
 }
 
-void DebugImage::draw_mask(const cv::Mat & mask, const cv::Scalar & color, double opacity)
-{
+void DebugImage::draw_mask(const cv::Mat& mask, const cv::Scalar& color, double opacity) {
   if (!active_ || mask.empty()) {
     return;
   }

src/bitbots_vision/src/model_config.cpp

@@ -1,14 +1,13 @@
 #include "bitbots_vision/model_config.hpp"
 
+#include <yaml-cpp/yaml.h>
+
 #include <filesystem>
 #include <stdexcept>
 
-#include <yaml-cpp/yaml.h>
-
 namespace bitbots_vision {
 
-ModelConfig ModelConfig::load_from(const std::string & model_dir)
-{
+ModelConfig ModelConfig::load_from(const std::string& model_dir) {
   const std::string path = model_dir + "/model_config.yaml";
   if (!std::filesystem::exists(path)) {
     throw std::runtime_error("model_config.yaml not found at: " + path);
@@ -18,11 +17,11 @@ ModelConfig ModelConfig::load_from(const std::string & model_dir)
 
   ModelConfig cfg;
 
-  for (const auto & name : root["detection"]["classes"]) {
+  for (const auto& name : root["detection"]["classes"]) {
     cfg.detection_classes_.push_back(name.as<std::string>());
   }
 
-  for (const auto & name : root["segmentation"]["classes"]) {
+  for (const auto& name : root["segmentation"]["classes"]) {
     cfg.segmentation_classes_.push_back(name.as<std::string>());
   }
 
@@ -33,8 +32,7 @@ ModelConfig ModelConfig::load_from(const std::string & model_dir)
   return cfg;
 }
 
-std::vector<int> ModelConfig::robot_class_ids() const
-{
+std::vector<int> ModelConfig::robot_class_ids() const {
   std::vector<int> ids;
   for (int i = 0; i < static_cast<int>(detection_classes_.size()); ++i) {
     if (detection_classes_[i].find("robot") != std::string::npos) {