odin_ros_driver1/include/reid_trt_extractor.hpp

#pragma once

#include <cstdint>
#include <deque>
#include <memory>
#include <string>
#include <vector>

#include <opencv2/core.hpp>

namespace odin_ros_driver {

using ReIDFeature = std::vector<float>;

struct ReIDTensorRTConfig {
    std::string engine_path;
    int input_height = 256;
    int input_width = 128;
    int feature_dim = 512;
    int max_batch_size = 8;
};

class TrackGallery {
public:
    explicit TrackGallery(size_t max_features = 10);

    void clear();
    void add(
        const ReIDFeature& feature,
        int64_t frame_index,
        const cv::Mat& crop_bgr = cv::Mat());

    bool empty() const { return features_.empty(); }
    size_t size() const { return features_.size(); }
    int64_t last_update_frame() const { return last_update_frame_; }
    const std::deque<cv::Mat>& crops_bgr() const { return crops_bgr_; }

    float best_similarity(const ReIDFeature& query) const;

private:
    std::deque<ReIDFeature> features_;
    std::deque<cv::Mat> crops_bgr_;
    size_t max_features_ = 10;
    int64_t last_update_frame_ = -1;
};

class ReIDTensorRTExtractor {
public:
    ReIDTensorRTExtractor();
    ~ReIDTensorRTExtractor();

    ReIDTensorRTExtractor(const ReIDTensorRTExtractor&) = delete;
    ReIDTensorRTExtractor& operator=(const ReIDTensorRTExtractor&) = delete;

    void initialize(const ReIDTensorRTConfig& config);
    bool initialized() const { return initialized_; }
    int feature_dim() const;

    std::vector<ReIDFeature> extract(
        const cv::Mat& frame_bgr,
        const std::vector<cv::Rect>& boxes) const;

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
    ReIDTensorRTConfig config_;
    bool initialized_ = false;
};

}  // namespace odin_ros_driver