Add AMD support via ort alternative to trt
This commit is contained in:
@@ -0,0 +1,177 @@
|
||||
// ── Gallery GEMM backend ──────────────────────────────────────────────────────
|
||||
// GPU similarity engine for the identity matcher. Uploads the reference gallery
|
||||
// once and computes the per-frame similarity matrix with a single SGEMM. The GPU
|
||||
// math library is selected at compile time by CMake (SAE_GEMM_BACKEND):
|
||||
// cuBLAS/CUDA or rocBLAS/HIP. This is the ONLY translation unit that includes
|
||||
// cublas/cuda or rocblas/hip headers.
|
||||
|
||||
#include "inference/similarity.hpp"
|
||||
|
||||
#if defined(SAE_GEMM_CUDA)
|
||||
#include <cublas_v2.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
#elif defined(SAE_GEMM_ROCM)
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <rocblas/rocblas.h>
|
||||
#else
|
||||
#error "gemm_backend.cpp requires SAE_GEMM_CUDA or SAE_GEMM_ROCM to be defined"
|
||||
#endif
|
||||
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
struct GpuError : std::runtime_error {
|
||||
using std::runtime_error::runtime_error;
|
||||
};
|
||||
|
||||
#if defined(SAE_GEMM_CUDA)
|
||||
|
||||
using stream_t = cudaStream_t;
|
||||
using blas_handle_t = cublasHandle_t;
|
||||
|
||||
inline void check_gpu(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess)
|
||||
throw GpuError(std::string(what) + ": " + cudaGetErrorString(e));
|
||||
}
|
||||
inline void check_blas(cublasStatus_t s, const char* what) {
|
||||
if (s != CUBLAS_STATUS_SUCCESS)
|
||||
throw GpuError(std::string(what) + ": cublas error " + std::to_string(s));
|
||||
}
|
||||
|
||||
inline void gpu_malloc(void** p, size_t bytes) { check_gpu(cudaMalloc(p, bytes), "cudaMalloc"); }
|
||||
inline void gpu_free(void* p) { cudaFree(p); }
|
||||
inline void gpu_memcpy_h2d(void* dst, const void* src, size_t n, stream_t s) { check_gpu(cudaMemcpyAsync(dst, src, n, cudaMemcpyHostToDevice, s), "H2D"); }
|
||||
inline void gpu_memcpy_d2h(void* dst, const void* src, size_t n, stream_t s) { check_gpu(cudaMemcpyAsync(dst, src, n, cudaMemcpyDeviceToHost, s), "D2H"); }
|
||||
inline void gpu_memcpy_h2d_sync(void* dst, const void* src, size_t n) { check_gpu(cudaMemcpy(dst, src, n, cudaMemcpyHostToDevice), "H2D_sync"); }
|
||||
inline void stream_create(stream_t* s) { check_gpu(cudaStreamCreate(s), "cudaStreamCreate"); }
|
||||
inline void stream_destroy(stream_t s) { cudaStreamDestroy(s); }
|
||||
inline void stream_sync(stream_t s) { check_gpu(cudaStreamSynchronize(s), "cudaStreamSync"); }
|
||||
inline void blas_create(blas_handle_t* h) { check_blas(cublasCreate(h), "cublasCreate"); }
|
||||
inline void blas_destroy(blas_handle_t h) { cublasDestroy(h); }
|
||||
inline void blas_set_stream(blas_handle_t h, stream_t s) { check_blas(cublasSetStream(h, s), "cublasSetStream"); }
|
||||
inline void blas_sgemm(blas_handle_t h, int m, int n, int k,
|
||||
const float* A, const float* B, float* C) {
|
||||
const float alpha = 1.f, beta = 0.f;
|
||||
check_blas(cublasSgemm(h, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
m, n, k, &alpha, A, k, B, k, &beta, C, m),
|
||||
"cublasSgemm");
|
||||
}
|
||||
inline const char* backend_name() { return "cuBLAS/CUDA"; }
|
||||
|
||||
#else // SAE_GEMM_ROCM
|
||||
|
||||
using stream_t = hipStream_t;
|
||||
using blas_handle_t = rocblas_handle;
|
||||
|
||||
inline void check_gpu(hipError_t e, const char* what) {
|
||||
if (e != hipSuccess)
|
||||
throw GpuError(std::string(what) + ": " + hipGetErrorString(e));
|
||||
}
|
||||
inline void check_blas(rocblas_status s, const char* what) {
|
||||
if (s != rocblas_status_success)
|
||||
throw GpuError(std::string(what) + ": rocblas error " + std::to_string(s));
|
||||
}
|
||||
|
||||
inline void gpu_malloc(void** p, size_t bytes) { check_gpu(hipMalloc(p, bytes), "hipMalloc"); }
|
||||
inline void gpu_free(void* p) { (void)hipFree(p); }
|
||||
inline void gpu_memcpy_h2d(void* dst, const void* src, size_t n, stream_t s) { check_gpu(hipMemcpyAsync(dst, src, n, hipMemcpyHostToDevice, s), "H2D"); }
|
||||
inline void gpu_memcpy_d2h(void* dst, const void* src, size_t n, stream_t s) { check_gpu(hipMemcpyAsync(dst, src, n, hipMemcpyDeviceToHost, s), "D2H"); }
|
||||
inline void gpu_memcpy_h2d_sync(void* dst, const void* src, size_t n) { check_gpu(hipMemcpy(dst, src, n, hipMemcpyHostToDevice), "H2D_sync"); }
|
||||
inline void stream_create(stream_t* s) { check_gpu(hipStreamCreate(s), "hipStreamCreate"); }
|
||||
inline void stream_destroy(stream_t s) { (void)hipStreamDestroy(s); }
|
||||
inline void stream_sync(stream_t s) { check_gpu(hipStreamSynchronize(s), "hipStreamSync"); }
|
||||
inline void blas_create(blas_handle_t* h) { check_blas(rocblas_create_handle(h), "rocblas_create_handle"); }
|
||||
inline void blas_destroy(blas_handle_t h) { rocblas_destroy_handle(h); }
|
||||
inline void blas_set_stream(blas_handle_t h, stream_t s) { check_blas(rocblas_set_stream(h, s), "rocblas_set_stream"); }
|
||||
inline void blas_sgemm(blas_handle_t h, int m, int n, int k,
|
||||
const float* A, const float* B, float* C) {
|
||||
const float alpha = 1.f, beta = 0.f;
|
||||
// rocblas_sgemm is column-major; same transposition trick as cuBLAS:
|
||||
// C(m×n) = A(k×m)^T * B(k×n) → S(N_gallery × n_faces) = G^T * Q
|
||||
check_blas(rocblas_sgemm(h, rocblas_operation_transpose, rocblas_operation_none,
|
||||
m, n, k, &alpha, A, k, B, k, &beta, C, m),
|
||||
"rocblas_sgemm");
|
||||
}
|
||||
inline const char* backend_name() { return "rocBLAS/HIP"; }
|
||||
|
||||
#endif
|
||||
|
||||
constexpr int kDim = 512;
|
||||
|
||||
class SimilarityEngine final : public ISimilarityEngine {
|
||||
public:
|
||||
SimilarityEngine(const float* gallery_row_major, int n_gallery, int max_faces)
|
||||
: n_gallery_(n_gallery), max_faces_(max_faces)
|
||||
{
|
||||
const size_t gallery_floats = static_cast<size_t>(n_gallery_) * kDim;
|
||||
gpu_malloc(reinterpret_cast<void**>(&d_gallery_), gallery_floats * sizeof(float));
|
||||
gpu_memcpy_h2d_sync(d_gallery_, gallery_row_major, gallery_floats * sizeof(float));
|
||||
|
||||
gpu_malloc(reinterpret_cast<void**>(&d_query_),
|
||||
static_cast<size_t>(max_faces_) * kDim * sizeof(float));
|
||||
gpu_malloc(reinterpret_cast<void**>(&d_sims_),
|
||||
static_cast<size_t>(max_faces_) * n_gallery_ * sizeof(float));
|
||||
|
||||
stream_create(&stream_);
|
||||
blas_create(&handle_);
|
||||
blas_set_stream(handle_, stream_);
|
||||
|
||||
host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);
|
||||
|
||||
std::cerr << "[similarity] " << backend_name() << " engine: gallery resident on GPU ("
|
||||
<< (gallery_floats * sizeof(float)) / (1024 * 1024) << " MiB)\n";
|
||||
}
|
||||
|
||||
~SimilarityEngine() override {
|
||||
if (d_gallery_) gpu_free(d_gallery_);
|
||||
if (d_query_) gpu_free(d_query_);
|
||||
if (d_sims_) gpu_free(d_sims_);
|
||||
if (handle_) blas_destroy(handle_);
|
||||
if (stream_) stream_destroy(stream_);
|
||||
}
|
||||
|
||||
SimilarityEngine(const SimilarityEngine&) = delete;
|
||||
SimilarityEngine& operator=(const SimilarityEngine&) = delete;
|
||||
|
||||
int max_faces() const override { return max_faces_; }
|
||||
|
||||
const float* compute(const float* query_row_major, int n_faces) override {
|
||||
if (n_faces <= 0) return host_sims_.data();
|
||||
if (n_faces > max_faces_)
|
||||
throw std::runtime_error("SimilarityEngine: n_faces exceeds max_faces");
|
||||
|
||||
gpu_memcpy_h2d(d_query_, query_row_major,
|
||||
static_cast<size_t>(n_faces) * kDim * sizeof(float), stream_);
|
||||
|
||||
// S (N_gallery × n_faces) col-major = G(512 × N_gallery)^T * Q(512 × n_faces)
|
||||
blas_sgemm(handle_, n_gallery_, n_faces, kDim, d_gallery_, d_query_, d_sims_);
|
||||
|
||||
gpu_memcpy_d2h(host_sims_.data(), d_sims_,
|
||||
static_cast<size_t>(n_gallery_) * n_faces * sizeof(float), stream_);
|
||||
stream_sync(stream_);
|
||||
return host_sims_.data();
|
||||
}
|
||||
|
||||
private:
|
||||
int n_gallery_{0};
|
||||
int max_faces_{0};
|
||||
float* d_gallery_{nullptr};
|
||||
float* d_query_{nullptr};
|
||||
float* d_sims_{nullptr};
|
||||
std::vector<float> host_sims_;
|
||||
stream_t stream_{};
|
||||
blas_handle_t handle_{};
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<ISimilarityEngine> make_similarity_engine(
|
||||
const float* gallery_row_major, int n_gallery, int max_faces) {
|
||||
return std::make_unique<SimilarityEngine>(gallery_row_major, n_gallery, max_faces);
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
// ── ORT inference backend ─────────────────────────────────────────────────────
|
||||
// ONNX Runtime 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=ORT.
|
||||
//
|
||||
// This is the ONLY translation unit that includes onnxruntime headers; the core
|
||||
// application never sees them.
|
||||
|
||||
#include "inference/face_detector.hpp"
|
||||
#include "inference/face_embedder.hpp"
|
||||
#include "backends/ort_provider.hpp"
|
||||
#include "config.hpp"
|
||||
#include "face_utils.hpp"
|
||||
#include "types.hpp"
|
||||
|
||||
#include <onnxruntime/onnxruntime_cxx_api.h>
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
// ── SCRFDDecoder ──────────────────────────────────────────────────────────────
|
||||
// ONNX Runtime SCRFD face detector with kps. Uses ORT (not cv::dnn) because
|
||||
// OpenCV 4.x cannot load SCRFD's dynamic Shape nodes. ORT handles dynamic shapes
|
||||
// natively and is thread-safe for concurrent Run() calls.
|
||||
//
|
||||
// Model output layout (9 tensors, InsightFace export order):
|
||||
// [0-2] score_s8 / score_s16 / score_s32 — flat (N,)
|
||||
// [3-5] bbox_s8 / bbox_s16 / bbox_s32 — flat (N*4,) distance format
|
||||
// [6-8] kps_s8 / kps_s16 / kps_s32 — flat (N*10,) distance format
|
||||
// Landmark order: right-eye, left-eye, nose, right-mouth, left-mouth.
|
||||
class SCRFDDecoder 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;
|
||||
|
||||
SCRFDDecoder(const std::string& model_path,
|
||||
float conf_threshold, float nms_threshold,
|
||||
OrtProvider provider, BackendConfig trt_cfg)
|
||||
: conf_threshold_(conf_threshold)
|
||||
, nms_threshold_(nms_threshold)
|
||||
{
|
||||
Ort::SessionOptions opts;
|
||||
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
opts.SetIntraOpNumThreads(1);
|
||||
|
||||
// SCRFD ONNX has a dynamic H/W input; we letterbox to 640×640 at
|
||||
// runtime, so pin the TRT-EP profile to that single shape.
|
||||
if (provider == OrtProvider::TensorRT) {
|
||||
if (trt_cfg.input_name.empty()) {
|
||||
Ort::SessionOptions probe_opts;
|
||||
probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
|
||||
Ort::Session probe(env_, model_path.c_str(), probe_opts);
|
||||
Ort::AllocatorWithDefaultOptions alloc;
|
||||
trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
|
||||
}
|
||||
if (trt_cfg.profile_min.empty()) {
|
||||
const std::string shape =
|
||||
"1x3x" + std::to_string(kInputH) + "x" + std::to_string(kInputW);
|
||||
trt_cfg.profile_min = shape;
|
||||
trt_cfg.profile_opt = shape;
|
||||
trt_cfg.profile_max = shape;
|
||||
}
|
||||
}
|
||||
|
||||
apply_ort_model_cache(opts, model_path, trt_cfg);
|
||||
apply_ort_provider(opts, provider, "SCRFDDecoder", trt_cfg);
|
||||
|
||||
session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);
|
||||
|
||||
Ort::AllocatorWithDefaultOptions alloc;
|
||||
auto in_name = session_->GetInputNameAllocated(0, alloc);
|
||||
input_name_ = in_name.get();
|
||||
|
||||
const size_t n_out = session_->GetOutputCount();
|
||||
if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
|
||||
throw std::runtime_error(
|
||||
"[SCRFDDecoder] expected 9 or 12 outputs (kps-variant model), got "
|
||||
+ std::to_string(n_out));
|
||||
fmc_ = static_cast<int>(n_out / 3);
|
||||
|
||||
for (size_t i = 0; i < n_out; ++i) {
|
||||
auto name = session_->GetOutputNameAllocated(i, alloc);
|
||||
out_name_storage_.emplace_back(name.get());
|
||||
}
|
||||
for (auto& s : out_name_storage_)
|
||||
out_name_ptrs_.push_back(s.c_str());
|
||||
|
||||
// Reject non-SCRFD models (e.g. YuNet, which also has 12 outputs).
|
||||
const int expected_last[3] = {1, 4, 10};
|
||||
for (size_t gi = 0; gi < 3; ++gi) {
|
||||
for (int si = 0; si < fmc_; ++si) {
|
||||
const size_t oi = gi * fmc_ + si;
|
||||
auto shape = session_->GetOutputTypeInfo(oi)
|
||||
.GetTensorTypeAndShapeInfo().GetShape();
|
||||
if (shape.empty() || shape.back() != expected_last[gi]) {
|
||||
throw std::runtime_error(
|
||||
"[SCRFDDecoder] model does not look like InsightFace SCRFD: "
|
||||
"output '" + out_name_storage_[oi] + "' last-dim is "
|
||||
+ std::to_string(shape.empty() ? -1 : shape.back())
|
||||
+ ", expected " + std::to_string(expected_last[gi])
|
||||
+ ". Hint: pass scrfd_500m_bnkps.onnx, not yunet/*.onnx.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::cerr << "[SCRFDDecoder] loaded: " << model_path << "\n";
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
const std::array<int64_t, 4> in_shape = {1, 3, kInputH, kInputW};
|
||||
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||||
auto in_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem, blob.ptr<float>(), blob.total(),
|
||||
in_shape.data(), in_shape.size());
|
||||
|
||||
const char* in_name_c = input_name_.c_str();
|
||||
auto outs = session_->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
&in_name_c, &in_tensor, 1,
|
||||
out_name_ptrs_.data(), out_name_ptrs_.size());
|
||||
|
||||
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 = outs[si].GetTensorData<float>();
|
||||
const float* b = outs[fmc_ + si].GetTensorData<float>();
|
||||
const float* k = outs[fmc_ * 2 + si].GetTensorData<float>();
|
||||
|
||||
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:
|
||||
float conf_threshold_;
|
||||
float nms_threshold_;
|
||||
int fmc_{3};
|
||||
Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "scrfd"};
|
||||
std::unique_ptr<Ort::Session> session_;
|
||||
std::string input_name_;
|
||||
std::vector<std::string> out_name_storage_;
|
||||
std::vector<const char*> out_name_ptrs_;
|
||||
};
|
||||
|
||||
// ── ArcFaceEmbedder ───────────────────────────────────────────────────────────
|
||||
// ONNX Runtime ArcFace embedder (w600k_r50, mbf, r18).
|
||||
// Input: [N, 3, 112, 112] float32, BGR→RGB, normalised to [-1, 1]
|
||||
// Output: [N, 512] float32 → L2-normalised per row
|
||||
class ArcFaceEmbedder final : public IFaceEmbedder {
|
||||
public:
|
||||
ArcFaceEmbedder(const std::string& model_path,
|
||||
OrtProvider provider, BackendConfig trt_cfg, int max_batch)
|
||||
: max_batch_(std::max(1, max_batch))
|
||||
{
|
||||
Ort::SessionOptions opts;
|
||||
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
opts.SetIntraOpNumThreads(1);
|
||||
|
||||
if (provider == OrtProvider::TensorRT) {
|
||||
if (trt_cfg.input_name.empty()) {
|
||||
Ort::SessionOptions probe_opts;
|
||||
probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
|
||||
Ort::Session probe(env_, model_path.c_str(), probe_opts);
|
||||
Ort::AllocatorWithDefaultOptions alloc;
|
||||
trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
|
||||
}
|
||||
if (trt_cfg.profile_min.empty()) {
|
||||
const std::string tail = "x3x112x112";
|
||||
trt_cfg.profile_min = "1" + tail;
|
||||
trt_cfg.profile_opt = std::to_string(max_batch_) + tail;
|
||||
trt_cfg.profile_max = std::to_string(max_batch_) + tail;
|
||||
}
|
||||
}
|
||||
|
||||
apply_ort_model_cache(opts, model_path, trt_cfg);
|
||||
apply_ort_provider(opts, provider, "ArcFace", trt_cfg);
|
||||
|
||||
session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);
|
||||
|
||||
Ort::AllocatorWithDefaultOptions alloc;
|
||||
auto in_name = session_->GetInputNameAllocated(0, alloc);
|
||||
auto out_name = session_->GetOutputNameAllocated(0, alloc);
|
||||
input_name_ = in_name.get();
|
||||
output_name_ = out_name.get();
|
||||
|
||||
auto in_type = session_->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetElementType();
|
||||
auto out_type = session_->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetElementType();
|
||||
input_is_fp16_ = (in_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
|
||||
output_is_fp16_ = (out_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
|
||||
|
||||
std::cerr << "[ArcFace] loaded: " << model_path << "\n";
|
||||
}
|
||||
|
||||
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());
|
||||
|
||||
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);
|
||||
|
||||
const std::array<int64_t, 4> in_shape = {n, 3, 112, 112};
|
||||
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
const char* in_name = input_name_.c_str();
|
||||
const char* out_name = output_name_.c_str();
|
||||
|
||||
cv::Mat blob16;
|
||||
if (input_is_fp16_) blob.convertTo(blob16, CV_16F);
|
||||
|
||||
Ort::Value in_tensor = input_is_fp16_
|
||||
? Ort::Value::CreateTensor<Ort::Float16_t>(
|
||||
mem, reinterpret_cast<Ort::Float16_t*>(blob16.ptr<uint16_t>()), blob16.total(),
|
||||
in_shape.data(), in_shape.size())
|
||||
: Ort::Value::CreateTensor<float>(
|
||||
mem, blob.ptr<float>(), blob.total(),
|
||||
in_shape.data(), in_shape.size());
|
||||
|
||||
auto outs = session_->Run(Ort::RunOptions{nullptr}, &in_name, &in_tensor, 1, &out_name, 1);
|
||||
|
||||
std::vector<Embedding> result(n);
|
||||
if (output_is_fp16_) {
|
||||
const auto* data16 = outs[0].GetTensorData<Ort::Float16_t>();
|
||||
std::vector<float> buf(n * 512);
|
||||
for (int j = 0; j < n * 512; ++j)
|
||||
buf[j] = data16[j].ToFloat();
|
||||
for (int i = 0; i < n; ++i)
|
||||
result[i] = l2_normalise(buf.data() + i * 512);
|
||||
} else {
|
||||
const float* data = outs[0].GetTensorData<float>();
|
||||
for (int i = 0; i < n; ++i)
|
||||
result[i] = l2_normalise(data + i * 512);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private:
|
||||
int max_batch_;
|
||||
Ort::Env env_{ORT_LOGGING_LEVEL_ERROR, "arcface"};
|
||||
std::unique_ptr<Ort::Session> session_;
|
||||
std::string input_name_;
|
||||
std::string output_name_;
|
||||
bool input_is_fp16_ = false;
|
||||
bool output_is_fp16_ = false;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ── Factories ─────────────────────────────────────────────────────────────────
|
||||
|
||||
std::unique_ptr<IFaceDetector> make_face_detector(const Config& cfg) {
|
||||
if (!cfg.detector_engine.empty())
|
||||
throw std::runtime_error(
|
||||
"detector_engine set but this build uses the ORT inference backend; "
|
||||
"rebuild with -DSAE_INFERENCE_BACKEND=TRT to use raw TensorRT engines.");
|
||||
const OrtProvider provider = detect_ort_provider();
|
||||
std::cerr << "[face_detector] ORT backend, provider: "
|
||||
<< provider_name(provider) << "\n";
|
||||
return std::make_unique<SCRFDDecoder>(
|
||||
cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
|
||||
}
|
||||
|
||||
std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg) {
|
||||
if (!cfg.arcface_engine.empty())
|
||||
throw std::runtime_error(
|
||||
"arcface_engine set but this build uses the ORT inference backend; "
|
||||
"rebuild with -DSAE_INFERENCE_BACKEND=TRT to use raw TensorRT engines.");
|
||||
const OrtProvider provider = detect_ort_provider();
|
||||
std::cerr << "[face_embedder] ORT backend, provider: "
|
||||
<< provider_name(provider) << "\n";
|
||||
return std::make_unique<ArcFaceEmbedder>(
|
||||
cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
#pragma once
|
||||
#include "inference/backend_config.hpp"
|
||||
|
||||
#include <onnxruntime/onnxruntime_cxx_api.h>
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
// Private to the ORT backend (backends/ort_backend.cpp). Detects the best
|
||||
// available ORT execution provider and applies it to a SessionOptions.
|
||||
// Priority order: TensorRT EP > CUDA > ROCm > CPU.
|
||||
//
|
||||
// Note: "TensorRT" here is ORT's TensorRT *execution provider*, distinct from
|
||||
// the raw-TensorRT backend (backends/trt_backend.cpp). This header never reaches
|
||||
// core translation units.
|
||||
//
|
||||
// Detection is conservative: GetAvailableProviders() confirms ORT was compiled
|
||||
// with the provider, then AppendExecutionProvider_* is attempted inside a
|
||||
// try/catch so a missing runtime library degrades gracefully to the next tier.
|
||||
|
||||
enum class OrtProvider { CPU, CUDA, ROCm, TensorRT };
|
||||
|
||||
inline OrtProvider detect_ort_provider() {
|
||||
auto available = Ort::GetAvailableProviders();
|
||||
for (const auto& p : available) {
|
||||
#ifdef SAE_ORT_WITH_TRT_EP
|
||||
if (p == "TensorrtExecutionProvider") return OrtProvider::TensorRT;
|
||||
if (p == "CUDAExecutionProvider") return OrtProvider::CUDA;
|
||||
#endif
|
||||
if (p == "ROCMExecutionProvider") return OrtProvider::ROCm;
|
||||
}
|
||||
return OrtProvider::CPU;
|
||||
}
|
||||
|
||||
inline const char* provider_name(OrtProvider p) {
|
||||
switch (p) {
|
||||
case OrtProvider::TensorRT: return "TensorRT-EP";
|
||||
case OrtProvider::CUDA: return "CUDA";
|
||||
case OrtProvider::ROCm: return "ROCm";
|
||||
default: return "CPU";
|
||||
}
|
||||
}
|
||||
|
||||
// If trt_cfg.ort_cache_dir is set, configure ORT to write/read a pre-optimized
|
||||
// .ort model for model_path. Must be called before AppendExecutionProvider_*.
|
||||
inline void apply_ort_model_cache(Ort::SessionOptions& opts,
|
||||
const std::string& model_path,
|
||||
const BackendConfig& trt_cfg) {
|
||||
if (trt_cfg.ort_cache_dir.empty()) return;
|
||||
std::filesystem::create_directories(trt_cfg.ort_cache_dir);
|
||||
const std::string stem =
|
||||
std::filesystem::path(model_path).stem().string();
|
||||
const std::string cache_path =
|
||||
trt_cfg.ort_cache_dir + "/" + stem + ".ort";
|
||||
opts.SetOptimizedModelFilePath(cache_path.c_str());
|
||||
}
|
||||
|
||||
// Apply the given provider to opts. Falls back to CPU on failure and returns the
|
||||
// provider that was actually applied.
|
||||
inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
|
||||
OrtProvider provider,
|
||||
const char* label,
|
||||
const BackendConfig& trt_cfg = {}) {
|
||||
#ifdef SAE_ORT_WITH_TRT_EP
|
||||
if (provider == OrtProvider::TensorRT) {
|
||||
try {
|
||||
std::filesystem::create_directories(trt_cfg.cache_dir);
|
||||
|
||||
std::unordered_map<std::string, std::string> kv = {
|
||||
{"device_id", "0"},
|
||||
{"trt_max_workspace_size", "2147483648"},
|
||||
{"trt_fp16_enable", trt_cfg.fp16 ? "1" : "0"},
|
||||
{"trt_int8_enable", trt_cfg.int8 ? "1" : "0"},
|
||||
{"trt_engine_cache_enable", "1"},
|
||||
{"trt_engine_cache_path", trt_cfg.cache_dir},
|
||||
};
|
||||
|
||||
if (!trt_cfg.input_name.empty() && !trt_cfg.profile_min.empty()) {
|
||||
kv["trt_profile_min_shapes"] =
|
||||
trt_cfg.input_name + ":" + trt_cfg.profile_min;
|
||||
kv["trt_profile_opt_shapes"] =
|
||||
trt_cfg.input_name + ":" + trt_cfg.profile_opt;
|
||||
kv["trt_profile_max_shapes"] =
|
||||
trt_cfg.input_name + ":" + trt_cfg.profile_max;
|
||||
}
|
||||
|
||||
Ort::TensorRTProviderOptions trt_v2;
|
||||
trt_v2.Update(kv);
|
||||
opts.AppendExecutionProvider_TensorRT_V2(*trt_v2);
|
||||
|
||||
std::cerr << "[" << label << "] TensorRT EP"
|
||||
<< (trt_cfg.fp16 ? " FP16" : "")
|
||||
<< (trt_cfg.int8 ? " INT8" : "")
|
||||
<< " cache=" << trt_cfg.cache_dir
|
||||
<< (trt_cfg.profile_min.empty() ? "" :
|
||||
" profile=" + trt_cfg.profile_min
|
||||
+ "/" + trt_cfg.profile_opt
|
||||
+ "/" + trt_cfg.profile_max)
|
||||
<< "\n";
|
||||
return OrtProvider::TensorRT;
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "[" << label << "] TensorRT EP unavailable ("
|
||||
<< e.what() << "), trying CUDA\n";
|
||||
provider = OrtProvider::CUDA;
|
||||
}
|
||||
}
|
||||
#endif // SAE_ORT_WITH_TRT_EP
|
||||
if (provider == OrtProvider::CUDA) {
|
||||
try {
|
||||
OrtCUDAProviderOptions cuda{};
|
||||
cuda.device_id = 0;
|
||||
opts.AppendExecutionProvider_CUDA(cuda);
|
||||
std::cerr << "[" << label << "] CUDA provider\n";
|
||||
return OrtProvider::CUDA;
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "[" << label << "] CUDA unavailable ("
|
||||
<< e.what() << "), trying ROCm\n";
|
||||
provider = OrtProvider::ROCm;
|
||||
}
|
||||
}
|
||||
if (provider == OrtProvider::ROCm) {
|
||||
try {
|
||||
OrtROCMProviderOptions rocm{};
|
||||
rocm.device_id = 0;
|
||||
opts.AppendExecutionProvider_ROCM(rocm);
|
||||
std::cerr << "[" << label << "] ROCm provider\n";
|
||||
return OrtProvider::ROCm;
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "[" << label << "] ROCm unavailable ("
|
||||
<< e.what() << "), falling back to CPU\n";
|
||||
}
|
||||
}
|
||||
std::cerr << "[" << label << "] CPU provider\n";
|
||||
return OrtProvider::CPU;
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
Reference in New Issue
Block a user