sam_seg_module.h

Go to the documentation of this file.

#pragma once
 
#include <opencv2/core.hpp>
 
#include <geolib/datatypes.h>
#include <ros/publisher.h>
#include <tue/config/configuration.h>
 
#include "ed/kinect/entity_update.h"
 
#include <atomic>
#include <memory>
#include <string>
#include <utility>
#include <vector>
 
struct SegmentationLatency
{
    double yolo_init_ms = 0.0;    // YOLO model initialization time
    double yolo_infer_ms = 0.0;   // YOLO inference time
    double sam_init_ms = 0.0;     // SAM model initialization time
    double sam_infer_ms = 0.0;    // SAM inference time
    double total_pipeline_ms = 0.0; // Total pipeline time
};
 
struct SegmentationResult
{
    std::vector<cv::Mat> masks;
    std::vector<cv::Rect> boxes;
    std::vector<std::string> labels;      // YOLO class name per detection
    std::vector<float> confidences;       // YOLO confidence per detection
    SegmentationLatency latency;          // Timing measurements (populated when verbose=true)
};
 
class SamSegPipeline
{
public:
    SamSegPipeline();
    ~SamSegPipeline();
 
    void initialize(tue::Configuration& config);
 
    SegmentationResult process(const cv::Mat& img, const cv::Mat& depth_image = cv::Mat(), const std::string& ignore_label = "", bool verbose = false);
 
private:
    struct Impl;
    std::unique_ptr<Impl> pimpl_;
    std::atomic<bool> is_initialized_;
};
 
void overlayMasksOnImage_(cv::Mat& rgb, const std::vector<cv::Mat>& masks);