376 lines
13 KiB
C++
376 lines
13 KiB
C++
#include "reid_trt_extractor.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <cstring>
|
|
#include <numeric>
|
|
#include <stdexcept>
|
|
#include <utility>
|
|
|
|
#include <cuda_fp16.h>
|
|
#include <opencv2/imgproc.hpp>
|
|
#include <yolos/core/trt_utils.hpp>
|
|
|
|
namespace odin_ros_driver {
|
|
|
|
namespace {
|
|
|
|
constexpr float kImageNetMean[3] = {0.485f, 0.456f, 0.406f};
|
|
constexpr float kImageNetStd[3] = {0.229f, 0.224f, 0.225f};
|
|
|
|
float l2_norm(const ReIDFeature& feature)
|
|
{
|
|
float value = 0.0f;
|
|
for (const float v : feature) {
|
|
value += v * v;
|
|
}
|
|
return std::sqrt(value);
|
|
}
|
|
|
|
void l2_normalize(ReIDFeature* feature)
|
|
{
|
|
if (feature == nullptr || feature->empty()) {
|
|
return;
|
|
}
|
|
|
|
const float norm = l2_norm(*feature);
|
|
if (norm <= 1e-6f) {
|
|
std::fill(feature->begin(), feature->end(), 0.0f);
|
|
return;
|
|
}
|
|
|
|
for (float& value : *feature) {
|
|
value /= norm;
|
|
}
|
|
}
|
|
|
|
nvinfer1::Dims make_dims4(int n, int c, int h, int w)
|
|
{
|
|
nvinfer1::Dims dims{};
|
|
dims.nbDims = 4;
|
|
dims.d[0] = n;
|
|
dims.d[1] = c;
|
|
dims.d[2] = h;
|
|
dims.d[3] = w;
|
|
return dims;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
TrackGallery::TrackGallery(size_t max_features)
|
|
: max_features_(std::max<size_t>(1, max_features))
|
|
{
|
|
}
|
|
|
|
void TrackGallery::clear()
|
|
{
|
|
features_.clear();
|
|
last_update_frame_ = -1;
|
|
}
|
|
|
|
void TrackGallery::add(const ReIDFeature& feature, int64_t frame_index)
|
|
{
|
|
if (feature.empty()) {
|
|
return;
|
|
}
|
|
|
|
if (features_.size() >= max_features_) {
|
|
features_.pop_front();
|
|
}
|
|
features_.push_back(feature);
|
|
last_update_frame_ = frame_index;
|
|
}
|
|
|
|
float TrackGallery::best_similarity(const ReIDFeature& query) const
|
|
{
|
|
if (query.empty() || features_.empty()) {
|
|
return -1.0f;
|
|
}
|
|
|
|
float best = -1.0f;
|
|
for (const auto& feature : features_) {
|
|
const size_t count = std::min(feature.size(), query.size());
|
|
float similarity = 0.0f;
|
|
for (size_t i = 0; i < count; ++i) {
|
|
similarity += feature[i] * query[i];
|
|
}
|
|
best = std::max(best, similarity);
|
|
}
|
|
return best;
|
|
}
|
|
|
|
struct ReIDTensorRTExtractor::Impl {
|
|
struct Binding {
|
|
std::string name;
|
|
bool is_input = false;
|
|
nvinfer1::DataType dtype = nvinfer1::DataType::kFLOAT;
|
|
size_t max_elements = 0;
|
|
yolos::trt::DeviceBuffer device;
|
|
};
|
|
|
|
yolos::trt::TrtUniquePtr<nvinfer1::IRuntime> runtime{nullptr};
|
|
yolos::trt::TrtUniquePtr<nvinfer1::ICudaEngine> engine{nullptr};
|
|
yolos::trt::TrtUniquePtr<nvinfer1::IExecutionContext> context{nullptr};
|
|
yolos::trt::CudaStream stream;
|
|
yolos::trt::PinnedBuffer host_input;
|
|
yolos::trt::PinnedBuffer host_output;
|
|
Binding input_binding;
|
|
Binding output_binding;
|
|
bool dynamic_batch = false;
|
|
int max_runtime_batch_size = 1;
|
|
};
|
|
|
|
ReIDTensorRTExtractor::ReIDTensorRTExtractor() = default;
|
|
|
|
ReIDTensorRTExtractor::~ReIDTensorRTExtractor() = default;
|
|
|
|
void ReIDTensorRTExtractor::initialize(const ReIDTensorRTConfig& config)
|
|
{
|
|
if (config.engine_path.empty()) {
|
|
throw std::runtime_error("ReID TensorRT engine path is empty");
|
|
}
|
|
|
|
auto impl = std::make_unique<Impl>();
|
|
const std::vector<char> engine_bytes = yolos::trt::readFile(config.engine_path);
|
|
|
|
impl->runtime.reset(nvinfer1::createInferRuntime(yolos::trt::getLogger()));
|
|
if (!impl->runtime) {
|
|
throw std::runtime_error("Failed to create TensorRT runtime for ReID");
|
|
}
|
|
|
|
impl->engine.reset(
|
|
impl->runtime->deserializeCudaEngine(engine_bytes.data(), engine_bytes.size()));
|
|
if (!impl->engine) {
|
|
throw std::runtime_error("Failed to deserialize ReID TensorRT engine: " + config.engine_path);
|
|
}
|
|
|
|
impl->context.reset(impl->engine->createExecutionContext());
|
|
if (!impl->context) {
|
|
throw std::runtime_error("Failed to create ReID TensorRT execution context");
|
|
}
|
|
|
|
if (impl->engine->getNbOptimizationProfiles() > 0 &&
|
|
!impl->context->setOptimizationProfileAsync(0, impl->stream.get())) {
|
|
throw std::runtime_error("Failed to select optimization profile 0 for ReID TensorRT engine");
|
|
}
|
|
|
|
bool found_input = false;
|
|
bool found_output = false;
|
|
for (int i = 0; i < impl->engine->getNbIOTensors(); ++i) {
|
|
const char* name = impl->engine->getIOTensorName(i);
|
|
if (name == nullptr) {
|
|
continue;
|
|
}
|
|
const auto mode = impl->engine->getTensorIOMode(name);
|
|
if (mode == nvinfer1::TensorIOMode::kINPUT && !found_input) {
|
|
impl->input_binding.name = name;
|
|
impl->input_binding.is_input = true;
|
|
impl->input_binding.dtype = impl->engine->getTensorDataType(name);
|
|
found_input = true;
|
|
} else if (mode == nvinfer1::TensorIOMode::kOUTPUT && !found_output) {
|
|
impl->output_binding.name = name;
|
|
impl->output_binding.is_input = false;
|
|
impl->output_binding.dtype = impl->engine->getTensorDataType(name);
|
|
found_output = true;
|
|
}
|
|
}
|
|
|
|
if (!found_input || !found_output) {
|
|
throw std::runtime_error("Expected one input and one output tensor in ReID TensorRT engine");
|
|
}
|
|
|
|
if (impl->input_binding.dtype != nvinfer1::DataType::kFLOAT) {
|
|
throw std::runtime_error("Only float32 ReID TensorRT inputs are supported");
|
|
}
|
|
if (impl->output_binding.dtype != nvinfer1::DataType::kFLOAT &&
|
|
impl->output_binding.dtype != nvinfer1::DataType::kHALF) {
|
|
throw std::runtime_error("Unsupported ReID TensorRT output dtype");
|
|
}
|
|
|
|
nvinfer1::Dims input_dims = impl->engine->getTensorShape(impl->input_binding.name.c_str());
|
|
if (input_dims.nbDims != 4) {
|
|
throw std::runtime_error("Expected ReID TensorRT input dims [N,C,H,W]");
|
|
}
|
|
|
|
config_ = config;
|
|
if (input_dims.d[2] > 0) {
|
|
config_.input_height = input_dims.d[2];
|
|
}
|
|
if (input_dims.d[3] > 0) {
|
|
config_.input_width = input_dims.d[3];
|
|
}
|
|
|
|
impl->dynamic_batch = (input_dims.d[0] == -1);
|
|
impl->max_runtime_batch_size = std::max(1, config_.max_batch_size);
|
|
|
|
nvinfer1::Dims resolved_input_dims = input_dims;
|
|
if (impl->dynamic_batch) {
|
|
const nvinfer1::Dims max_input_dims = make_dims4(
|
|
impl->max_runtime_batch_size, 3, config_.input_height, config_.input_width);
|
|
if (!impl->context->setInputShape(impl->input_binding.name.c_str(), max_input_dims)) {
|
|
throw std::runtime_error("Failed to set ReID TensorRT input shape");
|
|
}
|
|
resolved_input_dims = max_input_dims;
|
|
} else {
|
|
impl->max_runtime_batch_size = std::max(1, static_cast<int>(input_dims.d[0]));
|
|
}
|
|
|
|
const nvinfer1::Dims resolved_output_dims =
|
|
impl->context->getTensorShape(impl->output_binding.name.c_str());
|
|
if (resolved_output_dims.nbDims < 2) {
|
|
throw std::runtime_error("Unexpected ReID TensorRT output dimensions");
|
|
}
|
|
if (resolved_output_dims.d[resolved_output_dims.nbDims - 1] > 0) {
|
|
config_.feature_dim = resolved_output_dims.d[resolved_output_dims.nbDims - 1];
|
|
}
|
|
|
|
impl->input_binding.max_elements =
|
|
static_cast<size_t>(impl->max_runtime_batch_size) * 3U *
|
|
static_cast<size_t>(config_.input_height) * static_cast<size_t>(config_.input_width);
|
|
impl->output_binding.max_elements =
|
|
static_cast<size_t>(impl->max_runtime_batch_size) * static_cast<size_t>(config_.feature_dim);
|
|
|
|
impl->host_input.resize(impl->input_binding.max_elements * sizeof(float));
|
|
impl->host_output.resize(
|
|
impl->output_binding.max_elements * yolos::trt::dataTypeBytes(impl->output_binding.dtype));
|
|
impl->input_binding.device.resize(impl->host_input.size());
|
|
impl->output_binding.device.resize(impl->host_output.size());
|
|
|
|
if (!impl->context->setTensorAddress(
|
|
impl->input_binding.name.c_str(), impl->input_binding.device.data())) {
|
|
throw std::runtime_error("Failed to bind ReID TensorRT input tensor address");
|
|
}
|
|
if (!impl->context->setTensorAddress(
|
|
impl->output_binding.name.c_str(), impl->output_binding.device.data())) {
|
|
throw std::runtime_error("Failed to bind ReID TensorRT output tensor address");
|
|
}
|
|
|
|
impl_ = std::move(impl);
|
|
initialized_ = true;
|
|
}
|
|
|
|
int ReIDTensorRTExtractor::feature_dim() const
|
|
{
|
|
return config_.feature_dim;
|
|
}
|
|
|
|
std::vector<ReIDFeature> ReIDTensorRTExtractor::extract(
|
|
const cv::Mat& frame_bgr,
|
|
const std::vector<cv::Rect>& boxes) const
|
|
{
|
|
std::vector<ReIDFeature> features(
|
|
boxes.size(), ReIDFeature(static_cast<size_t>(config_.feature_dim), 0.0f));
|
|
if (!initialized_ || impl_ == nullptr || frame_bgr.empty() || boxes.empty()) {
|
|
return features;
|
|
}
|
|
|
|
std::vector<int> valid_indices;
|
|
std::vector<cv::Rect> valid_boxes;
|
|
valid_indices.reserve(boxes.size());
|
|
valid_boxes.reserve(boxes.size());
|
|
for (size_t i = 0; i < boxes.size(); ++i) {
|
|
const cv::Rect clamped = boxes[i] & cv::Rect(0, 0, frame_bgr.cols, frame_bgr.rows);
|
|
if (clamped.width < 4 || clamped.height < 4) {
|
|
continue;
|
|
}
|
|
valid_indices.push_back(static_cast<int>(i));
|
|
valid_boxes.push_back(clamped);
|
|
}
|
|
|
|
if (valid_boxes.empty()) {
|
|
return features;
|
|
}
|
|
|
|
const size_t pixels_per_crop =
|
|
static_cast<size_t>(config_.input_height) * static_cast<size_t>(config_.input_width);
|
|
const size_t floats_per_crop = 3U * pixels_per_crop;
|
|
auto* input_ptr = impl_->host_input.as<float>();
|
|
const size_t max_batch = static_cast<size_t>(std::max(1, impl_->max_runtime_batch_size));
|
|
|
|
for (size_t start = 0; start < valid_boxes.size(); start += max_batch) {
|
|
const size_t batch =
|
|
std::min(valid_boxes.size() - start, max_batch);
|
|
|
|
for (size_t local = 0; local < batch; ++local) {
|
|
const cv::Rect crop_box = valid_boxes[start + local];
|
|
cv::Mat resized;
|
|
cv::resize(
|
|
frame_bgr(crop_box),
|
|
resized,
|
|
cv::Size(config_.input_width, config_.input_height),
|
|
0.0,
|
|
0.0,
|
|
cv::INTER_LINEAR);
|
|
|
|
cv::Mat rgb;
|
|
cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB);
|
|
rgb.convertTo(rgb, CV_32FC3, 1.0 / 255.0);
|
|
|
|
std::vector<cv::Mat> channels(3);
|
|
cv::split(rgb, channels);
|
|
float* dst = input_ptr + local * floats_per_crop;
|
|
for (int channel = 0; channel < 3; ++channel) {
|
|
channels[channel] =
|
|
(channels[channel] - kImageNetMean[channel]) / kImageNetStd[channel];
|
|
std::memcpy(
|
|
dst + static_cast<size_t>(channel) * pixels_per_crop,
|
|
channels[channel].ptr<float>(),
|
|
pixels_per_crop * sizeof(float));
|
|
}
|
|
}
|
|
|
|
if (impl_->dynamic_batch) {
|
|
const nvinfer1::Dims batch_dims = make_dims4(
|
|
static_cast<int>(batch), 3, config_.input_height, config_.input_width);
|
|
if (!impl_->context->setInputShape(impl_->input_binding.name.c_str(), batch_dims)) {
|
|
throw std::runtime_error("Failed to set batch shape for ReID TensorRT inference");
|
|
}
|
|
}
|
|
|
|
const nvinfer1::Dims output_dims =
|
|
impl_->context->getTensorShape(impl_->output_binding.name.c_str());
|
|
const size_t output_elements = static_cast<size_t>(yolos::trt::volume(output_dims));
|
|
const size_t output_bytes =
|
|
output_elements * yolos::trt::dataTypeBytes(impl_->output_binding.dtype);
|
|
|
|
CUDA_CHECK(cudaMemcpyAsync(
|
|
impl_->input_binding.device.data(),
|
|
input_ptr,
|
|
batch * floats_per_crop * sizeof(float),
|
|
cudaMemcpyHostToDevice,
|
|
impl_->stream.get()));
|
|
|
|
if (!impl_->context->enqueueV3(impl_->stream.get())) {
|
|
throw std::runtime_error("TensorRT enqueueV3 failed for ReID inference");
|
|
}
|
|
|
|
CUDA_CHECK(cudaMemcpyAsync(
|
|
impl_->host_output.data(),
|
|
impl_->output_binding.device.data(),
|
|
output_bytes,
|
|
cudaMemcpyDeviceToHost,
|
|
impl_->stream.get()));
|
|
impl_->stream.synchronize();
|
|
|
|
for (size_t local = 0; local < batch; ++local) {
|
|
ReIDFeature feature(static_cast<size_t>(config_.feature_dim), 0.0f);
|
|
if (impl_->output_binding.dtype == nvinfer1::DataType::kFLOAT) {
|
|
const float* src = impl_->host_output.as<float>() + local * config_.feature_dim;
|
|
std::copy(src, src + config_.feature_dim, feature.begin());
|
|
} else {
|
|
const __half* src =
|
|
static_cast<const __half*>(impl_->host_output.data()) + local * config_.feature_dim;
|
|
for (int dim = 0; dim < config_.feature_dim; ++dim) {
|
|
feature[static_cast<size_t>(dim)] = __half2float(src[dim]);
|
|
}
|
|
}
|
|
l2_normalize(&feature);
|
|
features[static_cast<size_t>(valid_indices[start + local])] = std::move(feature);
|
|
}
|
|
}
|
|
|
|
return features;
|
|
}
|
|
|
|
} // namespace odin_ros_driver
|