dtourolle/scene-actor-extraction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
e1423062e260d6a393affc0c2ffe258dec7ce612
scene-actor-extraction/src/backends/trt_backend.cpp
T
dtourolle a5ee3c05ce feat(benchmark): per-node cost and bottleneck attribution for a run 
--benchmark <path> reports cumulative CPU and wall time per node and
names the node pacing the run. The pacing node is located from sampled
channel occupancy, not from time-in-node: backpressure inflates
time-in-node for everything downstream of the real bottleneck, so the
obvious measure names the victim rather than the cause.

Sampling starts with the network and stops before it is destroyed.
Channel fill is instantaneous and everything has drained by shutdown, so
a single read at the end reports an idle pipeline however congested it
was.

kill -USR1 dumps the table from a running or wedged process. Channel
occupancy identifies a stalled node -- full input, empty output --
without a debug build or a debugger, which is the difference between
diagnosing the AR-004 hang in seconds and reproducing it under gdb.

Two knobs this exposes for measurement rather than sets: SAE_CV_THREADS,
because OpenCV's TBB arena and KPN's thread-per-node are two schedulers
unaware of each other on the same cores; and SAE_CUDA_BLOCKING_SYNC,
because the default spin-wait held the embedder thread at 99.7% user
time while nvidia-powerd cut the GPU's clock from 1005 to 210 MHz.
Neither default changes until a measurement says it should.

TRACES: VR-015 | PR-004
2026-08-05 14:38:15 +02:00

626 lines
28 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// ── 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 <cstdlib>
#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));
}
/// TRACES: VR-015 | PR-004
/// Select how a thread waits for the GPU. Must run before the CUDA context is
/// created, so every engine constructor calls it and the first one wins.
///
/// The default (`cudaDeviceScheduleAuto`) spin-waits: `cudaStreamSynchronize`
/// burns the calling thread's CPU for the whole of the device's work. Measured
/// here, the embedder thread sat at 99.7% *user* time with 0.5 s of system time
/// across 183 s — i.e. no blocking syscalls at all — while the GPU ran flat out.
///
/// On this laptop that is not merely wasted CPU. `nvidia-powerd` arbitrates one
/// power budget across CPU and GPU, and the GPU's ceiling was observed dropping
/// from 20 W idle to 15 W under our load, with the SM clock *falling* from
/// 1005 MHz to 210 MHz once work started. Spinning may therefore be buying
/// watts away from the device the pipeline is actually waiting on.
///
/// SAE_CUDA_BLOCKING_SYNC=1 switches to a blocking wait so the A/B needs no
/// rebuild. Default is unchanged until the measurement says otherwise.
inline void configure_cuda_sync_once() {
static const bool done = [] {
const char* env = std::getenv("SAE_CUDA_BLOCKING_SYNC");
if (env && env[0] == '1') {
cudaError_t e = cudaSetDeviceFlags(cudaDeviceScheduleBlockingSync);
std::cerr << "[cuda] sync policy: BlockingSync"
<< (e == cudaSuccess ? "" : " (FAILED — context already created)")
<< "\n";
} else {
std::cerr << "[cuda] sync policy: default (spin) — "
"set SAE_CUDA_BLOCKING_SYNC=1 to compare\n";
}
return true;
}();
(void)done;
}
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");
configure_cuda_sync_once();
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;
}
configure_cuda_sync_once();
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");
configure_cuda_sync_once();
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);
}
Reference in New Issue View Git Blame Copy Permalink