OpenCV: distros (Arch/CachyOS) now ship OpenCV 5 as default. The config package rejects a 5.x install when find_package requests 4, so probe for 5 first and fall back to 4. All components used here (core, imgproc, imgcodecs, videoio, dnn, objdetect, highgui) exist in both. TensorRT: nvinfer1::Dims5 was removed in TRT 10 (Dims2..Dims4 remain in NvInferLegacyDims.h). Build the TransNetV2 rank-5 input shape via the generic nvinfer1::Dims, which is valid on both 8.x and 10.x.
588 lines
26 KiB
C++
588 lines
26 KiB
C++
// ── TensorRT inference backend ────────────────────────────────────────────────
|
||
// Pure-TensorRT implementations of IFaceDetector (SCRFD) and IFaceEmbedder
|
||
// (ArcFace), plus the make_* factories the core links against. Selected at
|
||
// compile time by CMake when SAE_INFERENCE_BACKEND=TRT.
|
||
//
|
||
// Loads serialised engines built by scripts/build_trt_engines.sh (or any
|
||
// trtexec-produced .engine matching the I/O contract). Skips ONNX Runtime
|
||
// entirely — useful where ORT was built without the TensorRT EP.
|
||
//
|
||
// This is the ONLY translation unit that includes NvInfer.h / cuda_runtime; the
|
||
// core application never sees them.
|
||
|
||
#include "inference/face_detector.hpp"
|
||
#include "inference/face_embedder.hpp"
|
||
#include "inference/scene_detector.hpp"
|
||
#include "config.hpp"
|
||
#include "face_utils.hpp"
|
||
#include "types.hpp"
|
||
|
||
#include <NvInfer.h>
|
||
#include <cuda_runtime_api.h>
|
||
#include <opencv2/dnn.hpp>
|
||
#include <opencv2/imgproc.hpp>
|
||
|
||
#include <algorithm>
|
||
#include <array>
|
||
#include <cmath>
|
||
#include <cstdint>
|
||
#include <fstream>
|
||
#include <iostream>
|
||
#include <memory>
|
||
#include <mutex>
|
||
#include <stdexcept>
|
||
#include <string>
|
||
#include <vector>
|
||
|
||
namespace {
|
||
|
||
struct CudaError : std::runtime_error {
|
||
using std::runtime_error::runtime_error;
|
||
};
|
||
|
||
inline void check_cuda(cudaError_t e, const char* what) {
|
||
if (e != cudaSuccess)
|
||
throw CudaError(std::string(what) + ": " + cudaGetErrorString(e));
|
||
}
|
||
|
||
class TrtLogger : public nvinfer1::ILogger {
|
||
public:
|
||
void log(Severity sev, const char* msg) noexcept override {
|
||
if (sev <= Severity::kWARNING)
|
||
std::cerr << "[TRT] " << msg << "\n";
|
||
}
|
||
};
|
||
inline TrtLogger& logger() { static TrtLogger g; return g; }
|
||
|
||
struct TrtDeleter { template<class T> void operator()(T* p) const { delete p; } };
|
||
|
||
inline std::vector<char> read_file(const std::string& path, const char* who) {
|
||
std::ifstream f(path, std::ios::binary | std::ios::ate);
|
||
if (!f) throw std::runtime_error(std::string(who) + ": cannot open " + path);
|
||
const std::streamsize sz = f.tellg();
|
||
f.seekg(0);
|
||
std::vector<char> blob(sz);
|
||
f.read(blob.data(), sz);
|
||
return blob;
|
||
}
|
||
|
||
// ── TrtArcFaceEmbedder ────────────────────────────────────────────────────────
|
||
// Engine I/O contract: input Nx3x112x112 float32/float16, output Nx512.
|
||
class TrtArcFaceEmbedder final : public IFaceEmbedder {
|
||
public:
|
||
explicit TrtArcFaceEmbedder(const std::string& engine_path) {
|
||
std::vector<char> blob = read_file(engine_path, "TrtArcFaceEmbedder");
|
||
|
||
runtime_.reset(nvinfer1::createInferRuntime(logger()));
|
||
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
||
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
|
||
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
|
||
context_.reset(engine_->createExecutionContext());
|
||
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
||
|
||
const int n_io = engine_->getNbIOTensors();
|
||
for (int i = 0; i < n_io; ++i) {
|
||
const char* name = engine_->getIOTensorName(i);
|
||
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)
|
||
input_name_ = name;
|
||
else
|
||
output_name_ = name;
|
||
}
|
||
if (input_name_.empty() || output_name_.empty())
|
||
throw std::runtime_error("TrtArcFaceEmbedder: engine missing input/output tensor");
|
||
|
||
auto in_dtype = engine_->getTensorDataType(input_name_.c_str());
|
||
auto out_dtype = engine_->getTensorDataType(output_name_.c_str());
|
||
input_is_fp16_ = (in_dtype == nvinfer1::DataType::kHALF);
|
||
output_is_fp16_ = (out_dtype == nvinfer1::DataType::kHALF);
|
||
|
||
auto max_dims = engine_->getProfileShape(input_name_.c_str(), 0,
|
||
nvinfer1::OptProfileSelector::kMAX);
|
||
if (max_dims.nbDims != 4 || max_dims.d[1] != 3 ||
|
||
max_dims.d[2] != 112 || max_dims.d[3] != 112)
|
||
throw std::runtime_error("TrtArcFaceEmbedder: unexpected input shape in engine");
|
||
max_batch_ = max_dims.d[0];
|
||
|
||
const std::size_t in_bytes = static_cast<std::size_t>(max_batch_) * 3 * 112 * 112 *
|
||
(input_is_fp16_ ? 2 : 4);
|
||
const std::size_t out_bytes = static_cast<std::size_t>(max_batch_) * 512 *
|
||
(output_is_fp16_ ? 2 : 4);
|
||
check_cuda(cudaMalloc(&d_input_, in_bytes), "cudaMalloc input");
|
||
check_cuda(cudaMalloc(&d_output_, out_bytes), "cudaMalloc output");
|
||
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
|
||
|
||
context_->setTensorAddress(input_name_.c_str(), d_input_);
|
||
context_->setTensorAddress(output_name_.c_str(), d_output_);
|
||
|
||
std::cerr << "[TrtArcFace] loaded: " << engine_path
|
||
<< " max_batch=" << max_batch_
|
||
<< (input_is_fp16_ ? " fp16-in" : "")
|
||
<< (output_is_fp16_ ? " fp16-out" : "")
|
||
<< "\n";
|
||
}
|
||
|
||
~TrtArcFaceEmbedder() override {
|
||
if (stream_) cudaStreamDestroy(stream_);
|
||
if (d_input_) cudaFree(d_input_);
|
||
if (d_output_) cudaFree(d_output_);
|
||
}
|
||
|
||
TrtArcFaceEmbedder(const TrtArcFaceEmbedder&) = delete;
|
||
TrtArcFaceEmbedder& operator=(const TrtArcFaceEmbedder&) = delete;
|
||
|
||
int max_batch() const override { return max_batch_; }
|
||
|
||
std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) override {
|
||
if (crops.empty()) return {};
|
||
const int n = static_cast<int>(crops.size());
|
||
if (n > max_batch_)
|
||
throw std::runtime_error("TrtArcFaceEmbedder: batch " + std::to_string(n) +
|
||
" exceeds engine max " + std::to_string(max_batch_));
|
||
|
||
std::vector<cv::Mat> rgbs(n);
|
||
for (int i = 0; i < n; ++i)
|
||
cv::cvtColor(crops[i], rgbs[i], cv::COLOR_BGR2RGB);
|
||
cv::Mat blob = cv::dnn::blobFromImages(
|
||
rgbs, 1.0 / 128.0, {112, 112},
|
||
cv::Scalar(127.5, 127.5, 127.5),
|
||
/*swapRB=*/false, /*crop=*/false, CV_32F);
|
||
|
||
std::lock_guard<std::mutex> lk(mu_);
|
||
context_->setInputShape(input_name_.c_str(),
|
||
nvinfer1::Dims4{n, 3, 112, 112});
|
||
|
||
const std::size_t in_count = static_cast<std::size_t>(n) * 3 * 112 * 112;
|
||
if (input_is_fp16_) {
|
||
cv::Mat blob16;
|
||
blob.convertTo(blob16, CV_16F);
|
||
check_cuda(cudaMemcpyAsync(d_input_, blob16.ptr(), in_count * 2,
|
||
cudaMemcpyHostToDevice, stream_),
|
||
"H2D input fp16");
|
||
} else {
|
||
check_cuda(cudaMemcpyAsync(d_input_, blob.ptr<float>(), in_count * 4,
|
||
cudaMemcpyHostToDevice, stream_),
|
||
"H2D input fp32");
|
||
}
|
||
|
||
if (!context_->enqueueV3(stream_))
|
||
throw std::runtime_error("TrtArcFaceEmbedder: enqueueV3 failed");
|
||
|
||
const std::size_t out_count = static_cast<std::size_t>(n) * 512;
|
||
std::vector<float> host_f32(out_count);
|
||
if (output_is_fp16_) {
|
||
std::vector<uint16_t> host_f16(out_count);
|
||
check_cuda(cudaMemcpyAsync(host_f16.data(), d_output_, out_count * 2,
|
||
cudaMemcpyDeviceToHost, stream_),
|
||
"D2H output fp16");
|
||
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||
cv::Mat src16(1, static_cast<int>(out_count), CV_16F, host_f16.data());
|
||
cv::Mat dst32(1, static_cast<int>(out_count), CV_32F, host_f32.data());
|
||
src16.convertTo(dst32, CV_32F);
|
||
} else {
|
||
check_cuda(cudaMemcpyAsync(host_f32.data(), d_output_, out_count * 4,
|
||
cudaMemcpyDeviceToHost, stream_),
|
||
"D2H output fp32");
|
||
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||
}
|
||
|
||
std::vector<Embedding> out(n);
|
||
for (int i = 0; i < n; ++i)
|
||
out[i] = l2_normalise(host_f32.data() + i * 512);
|
||
return out;
|
||
}
|
||
|
||
private:
|
||
std::unique_ptr<nvinfer1::IRuntime, TrtDeleter> runtime_;
|
||
std::unique_ptr<nvinfer1::ICudaEngine, TrtDeleter> engine_;
|
||
std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
|
||
|
||
std::string input_name_;
|
||
std::string output_name_;
|
||
bool input_is_fp16_ = false;
|
||
bool output_is_fp16_ = false;
|
||
int max_batch_ = 1;
|
||
|
||
void* d_input_ = nullptr;
|
||
void* d_output_ = nullptr;
|
||
cudaStream_t stream_ = nullptr;
|
||
|
||
mutable std::mutex mu_;
|
||
};
|
||
|
||
// ── TrtScrfdDecoder ───────────────────────────────────────────────────────────
|
||
// Pure-TensorRT SCRFD detector (1x3x640x640 input pinned). Post-processing
|
||
// matches the ORT decoder byte-for-byte — only inference is swapped.
|
||
class TrtScrfdDecoder final : public IFaceDetector {
|
||
public:
|
||
static constexpr int kInputW = 640;
|
||
static constexpr int kInputH = 640;
|
||
static constexpr int kAllStrides[4] = {8, 16, 32, 64};
|
||
static constexpr int kAnchors = 2;
|
||
|
||
TrtScrfdDecoder(const std::string& engine_path,
|
||
float conf_threshold, float nms_threshold)
|
||
: conf_threshold_(conf_threshold)
|
||
, nms_threshold_(nms_threshold)
|
||
{
|
||
std::vector<char> blob = read_file(engine_path, "TrtScrfdDecoder");
|
||
|
||
runtime_.reset(nvinfer1::createInferRuntime(logger()));
|
||
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
||
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
|
||
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
|
||
context_.reset(engine_->createExecutionContext());
|
||
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
||
|
||
const int n_io = engine_->getNbIOTensors();
|
||
for (int i = 0; i < n_io; ++i) {
|
||
const char* name = engine_->getIOTensorName(i);
|
||
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) {
|
||
if (!input_name_.empty())
|
||
throw std::runtime_error("TrtScrfdDecoder: multiple inputs not supported");
|
||
input_name_ = name;
|
||
} else {
|
||
output_names_.emplace_back(name);
|
||
}
|
||
}
|
||
if (input_name_.empty())
|
||
throw std::runtime_error("TrtScrfdDecoder: no input tensor");
|
||
const int n_out = static_cast<int>(output_names_.size());
|
||
if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
|
||
throw std::runtime_error(
|
||
"TrtScrfdDecoder: expected 9 or 12 outputs (kps-variant SCRFD), got "
|
||
+ std::to_string(n_out));
|
||
fmc_ = n_out / 3;
|
||
|
||
auto in_dims = engine_->getProfileShape(input_name_.c_str(), 0,
|
||
nvinfer1::OptProfileSelector::kOPT);
|
||
if (in_dims.nbDims != 4 || in_dims.d[0] != 1 || in_dims.d[1] != 3 ||
|
||
in_dims.d[2] != kInputH || in_dims.d[3] != kInputW)
|
||
throw std::runtime_error(
|
||
"TrtScrfdDecoder: engine input must be 1x3x" +
|
||
std::to_string(kInputH) + "x" + std::to_string(kInputW));
|
||
|
||
const std::size_t in_bytes = static_cast<std::size_t>(3) * kInputH * kInputW * 4;
|
||
check_cuda(cudaMalloc(&d_input_, in_bytes), "cudaMalloc input");
|
||
context_->setTensorAddress(input_name_.c_str(), d_input_);
|
||
context_->setInputShape(input_name_.c_str(),
|
||
nvinfer1::Dims4{1, 3, kInputH, kInputW});
|
||
|
||
d_outputs_.resize(n_out, nullptr);
|
||
host_outputs_.resize(n_out);
|
||
out_elem_counts_.resize(n_out, 0);
|
||
|
||
const int expected_last[3] = {1, 4, 10};
|
||
for (int oi = 0; oi < n_out; ++oi) {
|
||
auto dims = context_->getTensorShape(output_names_[oi].c_str());
|
||
if (dims.nbDims < 1)
|
||
throw std::runtime_error("TrtScrfdDecoder: bad shape for output " +
|
||
output_names_[oi]);
|
||
std::size_t count = 1;
|
||
for (int d = 0; d < dims.nbDims; ++d) count *= static_cast<std::size_t>(dims.d[d]);
|
||
const int last = dims.d[dims.nbDims - 1];
|
||
const int group = oi / fmc_; // 0=scores, 1=bboxes, 2=kps
|
||
if (last != expected_last[group])
|
||
throw std::runtime_error(
|
||
"TrtScrfdDecoder: output '" + output_names_[oi] + "' last-dim is " +
|
||
std::to_string(last) + ", expected " + std::to_string(expected_last[group]) +
|
||
". Engine does not match SCRFD-bnkps layout.");
|
||
|
||
check_cuda(cudaMalloc(&d_outputs_[oi], count * 4), "cudaMalloc output");
|
||
context_->setTensorAddress(output_names_[oi].c_str(), d_outputs_[oi]);
|
||
host_outputs_[oi].resize(count);
|
||
out_elem_counts_[oi] = count;
|
||
}
|
||
|
||
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
|
||
|
||
std::cerr << "[TrtScrfd] loaded: " << engine_path
|
||
<< " fmc=" << fmc_ << " outputs=" << n_out << "\n";
|
||
}
|
||
|
||
~TrtScrfdDecoder() override {
|
||
if (stream_) cudaStreamDestroy(stream_);
|
||
if (d_input_) cudaFree(d_input_);
|
||
for (void* p : d_outputs_) if (p) cudaFree(p);
|
||
}
|
||
|
||
TrtScrfdDecoder(const TrtScrfdDecoder&) = delete;
|
||
TrtScrfdDecoder& operator=(const TrtScrfdDecoder&) = delete;
|
||
|
||
std::vector<DetectedFace> detect(const cv::Mat& img) override {
|
||
const float scale = std::min(static_cast<float>(kInputW) / img.cols,
|
||
static_cast<float>(kInputH) / img.rows);
|
||
const int new_w = static_cast<int>(std::round(img.cols * scale));
|
||
const int new_h = static_cast<int>(std::round(img.rows * scale));
|
||
const int pad_x = (kInputW - new_w) / 2;
|
||
const int pad_y = (kInputH - new_h) / 2;
|
||
|
||
cv::Mat resized;
|
||
cv::resize(img, resized, {new_w, new_h}, 0, 0, cv::INTER_LINEAR);
|
||
cv::Mat letterboxed(kInputH, kInputW, img.type(), cv::Scalar(114, 114, 114));
|
||
resized.copyTo(letterboxed(cv::Rect(pad_x, pad_y, new_w, new_h)));
|
||
|
||
cv::Mat blob = cv::dnn::blobFromImage(
|
||
letterboxed, 1.0 / 128.0, {kInputW, kInputH},
|
||
cv::Scalar(127.5f, 127.5f, 127.5f),
|
||
/*swapRB=*/true, /*crop=*/false, CV_32F);
|
||
|
||
std::lock_guard<std::mutex> lk(mu_);
|
||
const std::size_t in_count = static_cast<std::size_t>(3) * kInputH * kInputW;
|
||
check_cuda(cudaMemcpyAsync(d_input_, blob.ptr<float>(), in_count * 4,
|
||
cudaMemcpyHostToDevice, stream_),
|
||
"H2D input");
|
||
|
||
if (!context_->enqueueV3(stream_))
|
||
throw std::runtime_error("TrtScrfdDecoder: enqueueV3 failed");
|
||
|
||
for (std::size_t oi = 0; oi < d_outputs_.size(); ++oi) {
|
||
check_cuda(cudaMemcpyAsync(host_outputs_[oi].data(), d_outputs_[oi],
|
||
out_elem_counts_[oi] * 4,
|
||
cudaMemcpyDeviceToHost, stream_),
|
||
"D2H output");
|
||
}
|
||
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||
|
||
std::vector<cv::Rect2d> raw_boxes;
|
||
std::vector<float> raw_scores;
|
||
std::vector<std::array<cv::Point2f, 5>> raw_kps;
|
||
|
||
for (int si = 0; si < fmc_; ++si) {
|
||
const int stride = kAllStrides[si];
|
||
const int fh = kInputH / stride;
|
||
const int fw = kInputW / stride;
|
||
|
||
const float* s = host_outputs_[si].data();
|
||
const float* b = host_outputs_[fmc_ + si].data();
|
||
const float* k = host_outputs_[fmc_ * 2 + si].data();
|
||
|
||
for (int r = 0; r < fh; ++r) {
|
||
for (int c = 0; c < fw; ++c) {
|
||
for (int a = 0; a < kAnchors; ++a) {
|
||
const int idx = (r * fw + c) * kAnchors + a;
|
||
const float score = s[idx];
|
||
if (score < conf_threshold_) continue;
|
||
|
||
const float cx = static_cast<float>(c * stride);
|
||
const float cy = static_cast<float>(r * stride);
|
||
|
||
const auto to_img_x = [&](float v) { return (v - pad_x) / scale; };
|
||
const auto to_img_y = [&](float v) { return (v - pad_y) / scale; };
|
||
|
||
const float x1 = to_img_x(cx - b[idx*4+0] * stride);
|
||
const float y1 = to_img_y(cy - b[idx*4+1] * stride);
|
||
const float x2 = to_img_x(cx + b[idx*4+2] * stride);
|
||
const float y2 = to_img_y(cy + b[idx*4+3] * stride);
|
||
raw_boxes.push_back({(double)x1, (double)y1,
|
||
(double)(x2-x1), (double)(y2-y1)});
|
||
raw_scores.push_back(score);
|
||
|
||
std::array<cv::Point2f, 5> lms;
|
||
for (int p = 0; p < 5; ++p)
|
||
lms[p] = {to_img_x(cx + k[idx*10+p*2 ] * stride),
|
||
to_img_y(cy + k[idx*10+p*2+1] * stride)};
|
||
raw_kps.push_back(lms);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
std::vector<int> keep;
|
||
cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep);
|
||
|
||
const float img_w = static_cast<float>(img.cols);
|
||
const float img_h = static_cast<float>(img.rows);
|
||
|
||
std::vector<DetectedFace> faces;
|
||
faces.reserve(keep.size());
|
||
for (int i : keep) {
|
||
const auto& rb = raw_boxes[i];
|
||
DetectedFace f;
|
||
const float x = std::max(0.f, (float)rb.x);
|
||
const float y = std::max(0.f, (float)rb.y);
|
||
f.bbox = {x, y,
|
||
std::min((float)rb.width, img_w - x),
|
||
std::min((float)rb.height, img_h - y)};
|
||
f.confidence = raw_scores[i];
|
||
f.landmarks = raw_kps[i];
|
||
faces.push_back(f);
|
||
}
|
||
return faces;
|
||
}
|
||
|
||
private:
|
||
std::unique_ptr<nvinfer1::IRuntime, TrtDeleter> runtime_;
|
||
std::unique_ptr<nvinfer1::ICudaEngine, TrtDeleter> engine_;
|
||
std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
|
||
|
||
float conf_threshold_;
|
||
float nms_threshold_;
|
||
int fmc_{3};
|
||
|
||
std::string input_name_;
|
||
std::vector<std::string> output_names_;
|
||
void* d_input_ = nullptr;
|
||
std::vector<void*> d_outputs_;
|
||
mutable std::vector<std::vector<float>> host_outputs_;
|
||
std::vector<std::size_t> out_elem_counts_;
|
||
|
||
cudaStream_t stream_ = nullptr;
|
||
mutable std::mutex mu_;
|
||
};
|
||
|
||
// ── TrtTransNetV2SceneDetector ────────────────────────────────────────────────
|
||
// Pure-TensorRT TransNetV2. Engine input pinned to 1×100×27×48×3 (RGB, 0-255);
|
||
// primary boundary head is the engine's first output tensor, sigmoided here to
|
||
// match the ORT path byte-for-byte.
|
||
class TrtTransNetV2SceneDetector final : public ISceneDetector {
|
||
public:
|
||
explicit TrtTransNetV2SceneDetector(const std::string& engine_path) {
|
||
std::vector<char> blob = read_file(engine_path, "TrtTransNetV2");
|
||
|
||
runtime_.reset(nvinfer1::createInferRuntime(logger()));
|
||
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
||
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
|
||
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
|
||
context_.reset(engine_->createExecutionContext());
|
||
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
||
|
||
const int n_io = engine_->getNbIOTensors();
|
||
for (int i = 0; i < n_io; ++i) {
|
||
const char* name = engine_->getIOTensorName(i);
|
||
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)
|
||
input_name_ = name;
|
||
else if (output_name_.empty())
|
||
output_name_ = name; // first output = primary boundary head
|
||
}
|
||
if (input_name_.empty() || output_name_.empty())
|
||
throw std::runtime_error("TrtTransNetV2: engine missing input/output tensor");
|
||
|
||
const std::size_t in_count =
|
||
static_cast<std::size_t>(kWindow) * kFrameH * kFrameW * 3;
|
||
const std::size_t out_count = static_cast<std::size_t>(kWindow);
|
||
check_cuda(cudaMalloc(&d_input_, in_count * 4), "cudaMalloc input");
|
||
check_cuda(cudaMalloc(&d_output_, out_count * 4), "cudaMalloc output");
|
||
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
|
||
context_->setTensorAddress(input_name_.c_str(), d_input_);
|
||
context_->setTensorAddress(output_name_.c_str(), d_output_);
|
||
|
||
std::cerr << "[TrtTransNetV2] loaded: " << engine_path
|
||
<< " (window=" << kWindow << ")\n";
|
||
}
|
||
|
||
~TrtTransNetV2SceneDetector() override {
|
||
if (stream_) cudaStreamDestroy(stream_);
|
||
if (d_input_) cudaFree(d_input_);
|
||
if (d_output_) cudaFree(d_output_);
|
||
}
|
||
|
||
TrtTransNetV2SceneDetector(const TrtTransNetV2SceneDetector&) = delete;
|
||
TrtTransNetV2SceneDetector& operator=(const TrtTransNetV2SceneDetector&) = delete;
|
||
|
||
std::vector<float> detect_window(const std::vector<cv::Mat>& window) override {
|
||
if (static_cast<int>(window.size()) != kWindow)
|
||
throw std::runtime_error(
|
||
"[TrtTransNetV2] detect_window expects exactly " +
|
||
std::to_string(kWindow) + " frames, got " +
|
||
std::to_string(window.size()));
|
||
|
||
std::vector<float> buf(
|
||
static_cast<size_t>(kWindow) * kFrameH * kFrameW * 3);
|
||
size_t o = 0;
|
||
for (int f = 0; f < kWindow; ++f) {
|
||
const cv::Mat& m = window[f];
|
||
const cv::Mat* src = &m;
|
||
cv::Mat resized;
|
||
if (m.cols != kFrameW || m.rows != kFrameH || m.type() != CV_8UC3) {
|
||
cv::Mat tmp;
|
||
if (m.type() != CV_8UC3) m.convertTo(tmp, CV_8UC3); else tmp = m;
|
||
cv::resize(tmp, resized, {kFrameW, kFrameH}, 0, 0, cv::INTER_AREA);
|
||
src = &resized;
|
||
}
|
||
for (int y = 0; y < kFrameH; ++y) {
|
||
const cv::Vec3b* row = src->ptr<cv::Vec3b>(y);
|
||
for (int x = 0; x < kFrameW; ++x) {
|
||
const cv::Vec3b& px = row[x]; // BGR
|
||
buf[o++] = static_cast<float>(px[2]); // R
|
||
buf[o++] = static_cast<float>(px[1]); // G
|
||
buf[o++] = static_cast<float>(px[0]); // B
|
||
}
|
||
}
|
||
}
|
||
|
||
std::lock_guard<std::mutex> lk(mu_);
|
||
// TensorRT 10 removed the fixed-rank Dims5 helper (Dims2..Dims4 remain
|
||
// in NvInferLegacyDims.h). Build the rank-5 shape via the generic Dims,
|
||
// which works on both 8.x and 10.x.
|
||
nvinfer1::Dims shape{};
|
||
shape.nbDims = 5;
|
||
shape.d[0] = 1;
|
||
shape.d[1] = kWindow;
|
||
shape.d[2] = kFrameH;
|
||
shape.d[3] = kFrameW;
|
||
shape.d[4] = 3;
|
||
context_->setInputShape(input_name_.c_str(), shape);
|
||
|
||
check_cuda(cudaMemcpyAsync(d_input_, buf.data(), buf.size() * 4,
|
||
cudaMemcpyHostToDevice, stream_), "H2D input");
|
||
if (!context_->enqueueV3(stream_))
|
||
throw std::runtime_error("TrtTransNetV2: enqueueV3 failed");
|
||
|
||
std::vector<float> logits(kWindow);
|
||
check_cuda(cudaMemcpyAsync(logits.data(), d_output_, kWindow * 4,
|
||
cudaMemcpyDeviceToHost, stream_), "D2H output");
|
||
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||
|
||
std::vector<float> probs(kWindow);
|
||
for (int i = 0; i < kWindow; ++i)
|
||
probs[i] = 1.f / (1.f + std::exp(-logits[i]));
|
||
return probs;
|
||
}
|
||
|
||
private:
|
||
std::unique_ptr<nvinfer1::IRuntime, TrtDeleter> runtime_;
|
||
std::unique_ptr<nvinfer1::ICudaEngine, TrtDeleter> engine_;
|
||
std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
|
||
std::string input_name_;
|
||
std::string output_name_;
|
||
void* d_input_ = nullptr;
|
||
void* d_output_ = nullptr;
|
||
cudaStream_t stream_ = nullptr;
|
||
mutable std::mutex mu_;
|
||
};
|
||
|
||
} // namespace
|
||
|
||
// ── Factories ─────────────────────────────────────────────────────────────────
|
||
|
||
std::unique_ptr<IFaceDetector> make_face_detector(const Config& cfg) {
|
||
if (cfg.detector_engine.empty())
|
||
throw std::runtime_error(
|
||
"TRT inference backend requires a pre-built detector engine "
|
||
"(--detector-engine / cfg.detector_engine). Build one with "
|
||
"scripts/build_trt_engines.sh, or rebuild with "
|
||
"-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
|
||
return std::make_unique<TrtScrfdDecoder>(
|
||
cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
|
||
}
|
||
|
||
std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg) {
|
||
if (cfg.arcface_engine.empty())
|
||
throw std::runtime_error(
|
||
"TRT inference backend requires a pre-built ArcFace engine "
|
||
"(--arcface-engine / cfg.arcface_engine). Build one with "
|
||
"scripts/build_trt_engines.sh, or rebuild with "
|
||
"-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
|
||
return std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
|
||
}
|
||
|
||
std::unique_ptr<ISceneDetector> make_scene_detector(const Config& cfg) {
|
||
if (cfg.scene_engine.empty())
|
||
throw std::runtime_error(
|
||
"TRT inference backend requires a pre-built TransNetV2 engine "
|
||
"(--scene-detector-engine / cfg.scene_engine). Build one with "
|
||
"scripts/convert_transnetv2.py --trt, or rebuild with "
|
||
"-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
|
||
return std::make_unique<TrtTransNetV2SceneDetector>(cfg.scene_engine);
|
||
}
|