71 lines
1.7 KiB
C++
71 lines
1.7 KiB
C++
#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
|