dtourolle/scene-actor-extraction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add AMD support via ort alternative to trt

Browse Source
This commit is contained in:
dtourolle
2026-06-28 11:50:05 +02:00
parent a3ba53ddf7
commit 0ee131a692
27 changed files with 1357 additions and 977 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backends/trt_backend.cpp
+455
View File
@@ -0,0 +1,455 @@
// ── 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 "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_;
};
} // 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);
}