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
.gitignore
+7 -1
View File
@@ -13,7 +13,7 @@ compile_commands.json
*.a
*.so
*.dylib
*.json
# Video files
*.mp4
*.mkv
@@ -27,6 +27,9 @@ external/*.onnx
# Generated TensorRT engines (rebuilt by ORT / scripts/build_trt_engines.sh)
trt_cache/
# ORT pre-optimized model cache (generated on first run, provider-specific)
ort_cache/
# Gallery JSON files (generated)
gallery.json
gallery_*.json
@@ -57,6 +60,9 @@ __pycache__/
.venv/
venv/
# Local secrets (API keys — never commit)
.env
# Editor / OS
.vscode/
.idea/
CMakeLists.txt
+162 -55
View File
@@ -21,46 +21,162 @@ add_subdirectory(external/KPN)
find_package(OpenCV 4 REQUIRED COMPONENTS
core imgproc imgcodecs videoio dnn objdetect highgui)
# ONNX Runtime (SCRFD detector — cv::dnn cannot handle dynamic Shape nodes)
find_library(ORT_LIB onnxruntime REQUIRED
HINTS /usr/lib /usr/local/lib)
find_path(ORT_INCLUDE onnxruntime_cxx_api.h
PATH_SUFFIXES onnxruntime
HINTS /usr/include /usr/local/include
REQUIRED)
add_library(onnxruntime UNKNOWN IMPORTED)
set_target_properties(onnxruntime PROPERTIES
IMPORTED_LOCATION "${ORT_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${ORT_INCLUDE}")
message(STATUS "ONNX Runtime: ${ORT_LIB} headers: ${ORT_INCLUDE}")
# ── Model paths ───────────────────────────────────────────────────────────────
# Defined early so the backend object libraries below can embed it.
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
CACHE PATH "Directory containing ONNX model files")
# TensorRT + CUDA runtime (raw-TRT ArcFace embedder; activated by --arcface-engine).
find_library(NVINFER_LIB nvinfer
HINTS /usr/lib /usr/local/lib /opt/tensorrt/lib)
find_path(NVINFER_INCLUDE NvInfer.h
HINTS /usr/include /usr/local/include /opt/tensorrt/include)
find_library(CUDART_LIB cudart
HINTS /opt/cuda/lib64 /usr/local/cuda/lib64 /usr/lib)
find_path(CUDART_INCLUDE cuda_runtime_api.h
HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include
/usr/local/cuda/include /usr/include)
find_library(CUBLAS_LIB cublas
HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64
/usr/local/cuda/lib64 /usr/lib)
if(NOT (NVINFER_LIB AND NVINFER_INCLUDE AND CUDART_LIB AND CUDART_INCLUDE AND CUBLAS_LIB))
message(FATAL_ERROR
"TensorRT or CUDA runtime not found "
"(nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE} "
"cudart=${CUDART_LIB} cublas=${CUBLAS_LIB} headers=${CUDART_INCLUDE})")
# ── Backend selection ─────────────────────────────────────────────────────────
# Two independent compile-time axes. The core application is agnostic to both:
# only the matching backend .cpp (in src/backends/) is compiled, and the backend
# headers (onnxruntime / NvInfer.h / cublas / rocblas) never reach core TUs.
#
# SAE_INFERENCE_BACKEND ORT → SCRFD + ArcFace via ONNX Runtime (.onnx models)
# TRT → SCRFD + ArcFace via raw TensorRT (.engine files)
# SAE_GEMM_BACKEND ROCM → gallery similarity GEMM via rocBLAS / HIP
# CUDA → gallery similarity GEMM via cuBLAS / CUDA
set(SAE_INFERENCE_BACKEND "ORT" CACHE STRING "Inference backend: ORT | TRT")
set(SAE_GEMM_BACKEND "ROCM" CACHE STRING "Gallery GEMM backend: ROCM | CUDA")
set_property(CACHE SAE_INFERENCE_BACKEND PROPERTY STRINGS ORT TRT)
set_property(CACHE SAE_GEMM_BACKEND PROPERTY STRINGS ROCM CUDA)
# Enable the ORT TensorRT/CUDA execution providers inside the ORT inference
# backend (only meaningful when ORT was built with the TensorRT EP). Off by
# default so ROCm/CPU builds don't reference unavailable EPs.
option(SAE_ORT_TRT_EP "ORT backend: enable TensorRT/CUDA execution providers" OFF)
# Back-compat: a legacy -DSAE_WITH_TRT=ON/OFF seeds the new vars (ON⇒TRT+CUDA,
# OFF⇒ORT+ROCM) unless the user set them explicitly.
if(DEFINED SAE_WITH_TRT)
if(SAE_WITH_TRT)
set(SAE_INFERENCE_BACKEND "TRT" CACHE STRING "" FORCE)
set(SAE_GEMM_BACKEND "CUDA" CACHE STRING "" FORCE)
else()
set(SAE_INFERENCE_BACKEND "ORT" CACHE STRING "" FORCE)
set(SAE_GEMM_BACKEND "ROCM" CACHE STRING "" FORCE)
endif()
message(STATUS "SAE_WITH_TRT=${SAE_WITH_TRT} (legacy) → "
"SAE_INFERENCE_BACKEND=${SAE_INFERENCE_BACKEND} "
"SAE_GEMM_BACKEND=${SAE_GEMM_BACKEND}")
endif()
add_library(trt_runtime INTERFACE)
target_include_directories(trt_runtime INTERFACE
"${NVINFER_INCLUDE}" "${CUDART_INCLUDE}")
target_link_libraries(trt_runtime INTERFACE
"${NVINFER_LIB}" "${CUDART_LIB}" "${CUBLAS_LIB}")
message(STATUS "TensorRT: ${NVINFER_LIB} CUDA runtime: ${CUDART_LIB} cuBLAS: ${CUBLAS_LIB}")
# FFmpeg (NVDEC hardware video decode + swscale colour conversion)
if(NOT SAE_INFERENCE_BACKEND MATCHES "^(ORT|TRT)$")
message(FATAL_ERROR "SAE_INFERENCE_BACKEND must be ORT or TRT (got '${SAE_INFERENCE_BACKEND}')")
endif()
if(NOT SAE_GEMM_BACKEND MATCHES "^(ROCM|CUDA)$")
message(FATAL_ERROR "SAE_GEMM_BACKEND must be ROCM or CUDA (got '${SAE_GEMM_BACKEND}')")
endif()
# CUDA runtime is needed by both TRT inference and CUDA GEMM — find it once.
function(sae_find_cudart)
if(TARGET cudart_dep)
return()
endif()
find_library(CUDART_LIB cudart
HINTS /opt/cuda/lib64 /usr/local/cuda/lib64 /usr/lib)
find_path(CUDART_INCLUDE cuda_runtime_api.h
HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include
/usr/local/cuda/include /usr/include)
if(NOT (CUDART_LIB AND CUDART_INCLUDE))
message(FATAL_ERROR "CUDA runtime not found (cudart=${CUDART_LIB} headers=${CUDART_INCLUDE}).")
endif()
add_library(cudart_dep INTERFACE)
target_include_directories(cudart_dep INTERFACE "${CUDART_INCLUDE}")
target_link_libraries(cudart_dep INTERFACE "${CUDART_LIB}")
set_property(GLOBAL PROPERTY sae_cudart_found TRUE)
endfunction()
# ── Inference backend dependency: builds the `inference_backend` object lib ────
if(SAE_INFERENCE_BACKEND STREQUAL "ORT")
find_library(ORT_LIB onnxruntime REQUIRED
HINTS /usr/lib64/rocm/lib /usr/lib /usr/local/lib)
find_path(ORT_INCLUDE onnxruntime_cxx_api.h
PATH_SUFFIXES onnxruntime
HINTS /usr/lib64/rocm/include/onnxruntime /usr/include/onnxruntime /usr/local/include/onnxruntime
/usr/lib64/rocm/include /usr/include /usr/local/include
REQUIRED)
# The include directive is <onnxruntime/onnxruntime_cxx_api.h>, so we need the
# parent of the onnxruntime/ subdirectory on the include path.
get_filename_component(ORT_INCLUDE_PARENT "${ORT_INCLUDE}" DIRECTORY)
if(NOT EXISTS "${ORT_INCLUDE_PARENT}/onnxruntime")
set(ORT_INCLUDE_PARENT "${ORT_INCLUDE}")
endif()
message(STATUS "Inference backend: ORT (${ORT_LIB} headers: ${ORT_INCLUDE_PARENT})")
add_library(inference_backend OBJECT src/backends/ort_backend.cpp)
set_target_properties(inference_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_include_directories(inference_backend PRIVATE src "${ORT_INCLUDE_PARENT}")
target_link_libraries(inference_backend PRIVATE ${OpenCV_LIBS} "${ORT_LIB}")
target_compile_definitions(inference_backend PRIVATE
SAE_MODELS_DIR="${SAE_MODELS_DIR}"
$<$<BOOL:${SAE_ORT_TRT_EP}>:SAE_ORT_WITH_TRT_EP>)
else() # TRT
find_library(NVINFER_LIB nvinfer
HINTS /usr/lib /usr/local/lib /opt/tensorrt/lib)
find_path(NVINFER_INCLUDE NvInfer.h
HINTS /usr/include /usr/local/include /opt/tensorrt/include)
if(NOT (NVINFER_LIB AND NVINFER_INCLUDE))
message(FATAL_ERROR
"TensorRT not found (nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE}). "
"Pass -DSAE_INFERENCE_BACKEND=ORT to load .onnx models without TensorRT.")
endif()
sae_find_cudart()
message(STATUS "Inference backend: TRT (${NVINFER_LIB})")
add_library(inference_backend OBJECT src/backends/trt_backend.cpp)
set_target_properties(inference_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_include_directories(inference_backend PRIVATE src "${NVINFER_INCLUDE}")
target_link_libraries(inference_backend PRIVATE
${OpenCV_LIBS} "${NVINFER_LIB}" cudart_dep)
target_compile_definitions(inference_backend PRIVATE
SAE_MODELS_DIR="${SAE_MODELS_DIR}")
endif()
# ── GEMM backend dependency: builds the `gemm_backend` object lib ──────────────
if(SAE_GEMM_BACKEND STREQUAL "CUDA")
find_library(CUBLAS_LIB cublas
HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64
/usr/local/cuda/lib64 /usr/lib)
if(NOT CUBLAS_LIB)
message(FATAL_ERROR "cuBLAS not found (cublas=${CUBLAS_LIB}).")
endif()
sae_find_cudart()
message(STATUS "GEMM backend: CUDA (${CUBLAS_LIB})")
add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_include_directories(gemm_backend PRIVATE src)
target_link_libraries(gemm_backend PRIVATE "${CUBLAS_LIB}" cudart_dep)
target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_CUDA)
else() # ROCM
find_library(ROCBLAS_LIB rocblas
HINTS /usr/lib64/rocm/lib /usr/lib64 /usr/local/lib)
find_path(ROCBLAS_INCLUDE rocblas/rocblas.h
HINTS /usr/lib64/rocm/include /usr/include /usr/local/include)
find_library(HIP_LIB amdhip64
HINTS /usr/lib64/rocm/lib /usr/lib64 /usr/local/lib)
find_path(HIP_INCLUDE hip/hip_runtime_api.h
HINTS /usr/lib64/rocm/include /usr/include /usr/local/include)
if(NOT (ROCBLAS_LIB AND ROCBLAS_INCLUDE AND HIP_LIB AND HIP_INCLUDE))
message(FATAL_ERROR
"rocBLAS or HIP runtime not found "
"(rocblas=${ROCBLAS_LIB} headers=${ROCBLAS_INCLUDE} "
"hip=${HIP_LIB} headers=${HIP_INCLUDE}). "
"Install rocblas-devel and hip-devel (or pass -DSAE_GEMM_BACKEND=CUDA).")
endif()
message(STATUS "GEMM backend: ROCM (${ROCBLAS_LIB})")
add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_include_directories(gemm_backend PRIVATE src "${ROCBLAS_INCLUDE}" "${HIP_INCLUDE}")
target_link_libraries(gemm_backend PRIVATE "${ROCBLAS_LIB}" "${HIP_LIB}")
# HIP headers require the platform to be declared explicitly when compiled with g++.
target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_ROCM __HIP_PLATFORM_AMD__)
endif()
# FFmpeg (hwaccel video decode: CUDA/VAAPI, runtime-detected + swscale colour
# conversion). Hwaccel support is built into libavcodec/libavutil; no extra
# libraries are needed here.
find_package(PkgConfig REQUIRED)
pkg_check_modules(AVFORMAT REQUIRED libavformat)
pkg_check_modules(AVCODEC REQUIRED libavcodec)
@@ -103,18 +219,22 @@ FetchContent_MakeAvailable(nanobind)
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
CACHE PATH "Directory containing ONNX model files")
# ── Shared library: gallery store ─────────────────────────────────────────────
# ── Shared library: gallery store + compiled-in backends ──────────────────────
# The backend object libraries carry their own ORT/TRT/CUDA/ROCm linkage and
# headers; sae_gallery re-exports those object files so every binary that links
# sae_gallery gets the chosen backend without ever seeing its headers.
add_library(sae_gallery STATIC
src/gallery/gallery_store.cpp
src/gallery/gallery_builder.cpp
)
set_target_properties(sae_gallery PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_include_directories(sae_gallery PUBLIC src)
target_link_libraries(sae_gallery PUBLIC
kpn
${OpenCV_LIBS}
nlohmann_json::nlohmann_json
onnxruntime
trt_runtime
inference_backend
gemm_backend
ffmpeg_libs
)
target_compile_definitions(sae_gallery PUBLIC
@@ -123,24 +243,11 @@ target_compile_definitions(sae_gallery PUBLIC
# ── embed_faces — image → embedding JSON (used by gallery builder scripts) ────
add_executable(embed_faces src/embed_faces.cpp)
target_link_libraries(embed_faces PRIVATE
kpn
${OpenCV_LIBS}
nlohmann_json::nlohmann_json
onnxruntime
trt_runtime
ffmpeg_libs
)
target_compile_definitions(embed_faces PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")
target_link_libraries(embed_faces PRIVATE sae_gallery)
# ── sae_embed — Python module: load SCRFD+ArcFace once, embed many images ───
nanobind_add_module(sae_embed src/python_bindings.cpp)
target_include_directories(sae_embed PRIVATE src)
target_link_libraries(sae_embed PRIVATE
${OpenCV_LIBS}
onnxruntime
)
target_compile_definitions(sae_embed PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")
target_link_libraries(sae_embed PRIVATE sae_gallery)
# ── analyze — main analysis binary ───────────────────────────────────────────
add_executable(scene_analyze src/main.cpp)
models/LVFace-B_Glint360K.onnx
LFS
BIN
View File
Binary file not shown.
scripts/make_jellyfin_gallery.py
+10 -6
View File
@@ -42,6 +42,7 @@ Get a free TMDB API key at: https://www.themoviedb.org/settings/api
import argparse
import concurrent.futures
import os
import io
import json
import sys
@@ -376,11 +377,13 @@ def main():
description="Build a gallery.json spanning an entire Jellyfin library",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--jellyfin-url", required=True,
parser.add_argument("--jellyfin-url", default=os.environ.get("JELLYFIN_URL"),
required=not os.environ.get("JELLYFIN_URL"),
help="Jellyfin base server URL only, e.g. http://jellyfin.local:8096 "
"(no /Items or other API path)")
parser.add_argument("--api-key", required=True,
help="Jellyfin API key (Dashboard → Advanced → API Keys)")
"(no /Items or other API path). Env: JELLYFIN_URL")
parser.add_argument("--api-key", default=os.environ.get("JELLYFIN_API_KEY"),
required=not os.environ.get("JELLYFIN_API_KEY"),
help="Jellyfin API key (Dashboard → Advanced → API Keys). Env: JELLYFIN_API_KEY")
parser.add_argument("--output", required=True, help="Output gallery.json path")
parser.add_argument("--item-types", default="Movie,Series",
help="Comma-separated Jellyfin item types to scan (default: Movie,Series)")
@@ -397,10 +400,11 @@ def main():
parser.add_argument("--fetch-imdb-ids", action="store_true",
help="Resolve each actor's real IMDB id via Jellyfin ProviderIds "
"(one extra API call per new actor; otherwise imdb_id is left empty)")
parser.add_argument("--tmdb-key", default=None,
parser.add_argument("--tmdb-key", default=os.environ.get("TMDB_API_KEY"),
help="TMDB API key/bearer token. If set, actors with no usable "
"Jellyfin image fall back to TMDB profile images (looked up "
"via the actor's IMDB id, requires one extra Jellyfin call per actor)")
"via the actor's IMDB id, requires one extra Jellyfin call per actor). "
"Env: TMDB_API_KEY")
parser.add_argument("--merge", action="store_true",
help="If --output already exists, keep its actors and only embed "
"actors not already present (matched by jellyfin_id)")
scripts/run_from_jellyfin.py
+7 -2
View File
@@ -29,6 +29,7 @@ random batch of items with no truth data yet, processing each in turn:
import argparse
import json
import os
import subprocess
import sys
import tempfile
@@ -174,8 +175,12 @@ def main():
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--jellyfin-url", required=True)
parser.add_argument("--api-key", required=True)
parser.add_argument("--jellyfin-url", default=os.environ.get("JELLYFIN_URL"),
required=not os.environ.get("JELLYFIN_URL"),
help="Jellyfin base URL. Env: JELLYFIN_URL")
parser.add_argument("--api-key", default=os.environ.get("JELLYFIN_API_KEY"),
required=not os.environ.get("JELLYFIN_API_KEY"),
help="Jellyfin API key. Env: JELLYFIN_API_KEY")
group = parser.add_mutually_exclusive_group()
group.add_argument("--item-id", help="Jellyfin item id of the title")
group.add_argument("--title", help="Title to search for (uses first match)")
src/arcface_embedder.hpp
-132
View File
@@ -1,132 +0,0 @@
#pragma once
#include "ort_provider.hpp"
#include "types.hpp"
#include "face_utils.hpp"
#include <onnxruntime/onnxruntime_cxx_api.h>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>
// ── ArcFaceEmbedder ───────────────────────────────────────────────────────────
// ONNX Runtime-based embedder for InsightFace ArcFace (w600k_r50, mbf, r18).
// Replaces cv::dnn::Net which has no GPU path and is ~5–10× slower.
//
// Input: [N, 3, 112, 112] float32, BGR→RGB, normalised to [-1, 1]
// Output: [N, 512] float32 → L2-normalised per row
//
// ORT Run() is thread-safe; no external locking is needed.
struct ArcFaceEmbedder {
explicit ArcFaceEmbedder(const std::string& model_path,
OrtProvider provider = OrtProvider::CPU,
TrtConfig trt_cfg = {},
int max_batch = 4)
{
Ort::SessionOptions opts;
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
opts.SetIntraOpNumThreads(1);
if (provider == OrtProvider::TensorRT) {
if (trt_cfg.input_name.empty()) {
// Probe input name from a tiny CPU session so we can configure
// the dynamic-batch profile before the real session is built.
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";
const int m = std::max(1, max_batch);
trt_cfg.profile_min = "1" + tail;
trt_cfg.profile_opt = std::to_string(m) + tail;
trt_cfg.profile_max = std::to_string(m) + tail;
}
}
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";
}
// Embed a batch of 112×112 BGR crops. Returns L2-normalised 512-d embeddings.
std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) const {
if (crops.empty()) return {};
const int n = static_cast<int>(crops.size());
// BGR→RGB, build NCHW float32 blob normalised to [-1, 1]
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;
}
Embedding embed_one(const cv::Mat& crop) const {
return embed({crop})[0];
}
private:
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;
};
src/backends/gemm_backend.cpp
+177
View File
@@ -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);
}
src/scrfd_decoder.hpp → src/backends/ort_backend.cpp
+167 -32
View File
@@ -1,31 +1,45 @@
#pragma once
#include "ort_provider.hpp"
// ── 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-based decoder for InsightFace SCRFD face detector with kps.
// Uses ORT instead of 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.
// 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 (same as YuNet/ArcFace convention):
// [0] right-eye [1] left-eye [2] nose [3] right-mouth [4] left-mouth
struct SCRFDDecoder {
// 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};
@@ -33,8 +47,7 @@ struct SCRFDDecoder {
SCRFDDecoder(const std::string& model_path,
float conf_threshold, float nms_threshold,
OrtProvider provider = OrtProvider::CPU,
TrtConfig trt_cfg = {})
OrtProvider provider, BackendConfig trt_cfg)
: conf_threshold_(conf_threshold)
, nms_threshold_(nms_threshold)
{
@@ -43,9 +56,7 @@ struct SCRFDDecoder {
opts.SetIntraOpNumThreads(1);
// SCRFD ONNX has a dynamic H/W input; we letterbox to 640×640 at
// runtime, so pin the TRT profile to that single shape — otherwise
// TRT picks generic shapes and either rebuilds per-call or falls
// back to CUDA EP.
// 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;
@@ -63,6 +74,7 @@ struct SCRFDDecoder {
}
}
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);
@@ -85,11 +97,7 @@ struct SCRFDDecoder {
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 and so
// passes the count check above, but is encoded entirely differently).
// Cross-check by output channel count: SCRFD's three groups of fmc_
// outputs encode scores (1ch), bboxes (4ch) and 5-point kps (10ch).
// YuNet exports loc/conf/iou with 14/2/1 channels, so this trips.
// 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) {
@@ -110,14 +118,7 @@ struct SCRFDDecoder {
std::cerr << "[SCRFDDecoder] loaded: " << model_path << "\n";
}
// Thread-safe: ORT Run() is safe for concurrent calls on the same Session.
std::vector<DetectedFace> detect(const cv::Mat& img) const {
// Letterbox to 640×640: uniform scale (preserves aspect ratio) + pad
// shorter side with constant grey. Stretching to 640×640 (the prior
// behaviour) distorts faces non-uniformly and degrades landmark
// localisation — matters most for portrait gallery images and 16:9
// video frames alike. Coordinates are mapped back via inverse scale +
// pad-offset below.
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));
@@ -131,7 +132,6 @@ struct SCRFDDecoder {
cv::Scalar(114, 114, 114));
resized.copyTo(letterboxed(cv::Rect(pad_x, pad_y, new_w, new_h)));
// BGR→RGB swap + normalize to [-1,1] → NCHW float32 blob
cv::Mat blob = cv::dnn::blobFromImage(
letterboxed, 1.0 / 128.0, {kInputW, kInputH},
cv::Scalar(127.5f, 127.5f, 127.5f),
@@ -172,8 +172,6 @@ struct SCRFDDecoder {
const float cx = static_cast<float>(c * stride);
const float cy = static_cast<float>(r * stride);
// Decode in letterboxed network space, then un-pad +
// un-scale to original image coordinates.
const auto to_img_x = [&](float v) { return (v - pad_x) / scale; };
const auto to_img_y = [&](float v) { return (v - pad_y) / scale; };
@@ -228,3 +226,140 @@ private:
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);
}
src/ort_provider.hpp → src/backends/ort_provider.hpp
+33 -42
View File
@@ -1,4 +1,6 @@
#pragma once
#include "inference/backend_config.hpp"
#include <onnxruntime/onnxruntime_cxx_api.h>
#include <filesystem>
#include <iostream>
@@ -6,8 +8,13 @@
#include <unordered_map>
#include <vector>
// Detect the best available ORT execution provider and apply it to a
// SessionOptions. Priority order: TensorRT > CUDA > ROCm > CPU.
// 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
@@ -15,36 +22,13 @@
enum class OrtProvider { CPU, CUDA, ROCm, TensorRT };
// ── TRT configuration ─────────────────────────────────────────────────────────
// fp16: FP16 Tensor Core kernels — safe for both SCRFD and ArcFace.
// int8: INT8 quantisation — fast but UNSAFE for ArcFace without a
// calibration table (embedding cosine space will shift, breaking
// your similarity thresholds). Safe for the SCRFD detector.
// cache_dir: TRT engines are compiled once and cached here. First run is
// slow (~30–60 s per model); every subsequent run loads instantly.
// Shape profile (optional, set input_name + profile_{min,opt,max} to enable):
// ArcFace input is dynamic-batch (Nx3x112x112) — without a profile TRT
// builds at batch=1 and any larger call falls back to CUDA EP.
// SCRFD input is static-batch with dynamic H/W; we letterbox to 640×640
// and pin the profile to that.
// Shape strings are trtexec-style, e.g. "1x3x112x112".
struct TrtConfig {
bool fp16 = true;
bool int8 = false;
std::string cache_dir = "./trt_cache";
// Per-tensor optimisation profile. All four fields must be set together.
std::string input_name; // e.g. "input.1"
std::string profile_min; // e.g. "1x3x112x112"
std::string profile_opt; // e.g. "4x3x112x112"
std::string profile_max; // e.g. "8x3x112x112"
};
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;
@@ -52,25 +36,38 @@ inline OrtProvider detect_ort_provider() {
inline const char* provider_name(OrtProvider p) {
switch (p) {
case OrtProvider::TensorRT: return "TensorRT";
case OrtProvider::TensorRT: return "TensorRT-EP";
case OrtProvider::CUDA: return "CUDA";
case OrtProvider::ROCm: return "ROCm";
default: return "CPU";
}
}
// Apply the given provider to opts. Falls back to CPU on failure and
// returns the provider that was actually applied.
// 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 TrtConfig& trt_cfg = {}) {
const BackendConfig& trt_cfg = {}) {
#ifdef SAE_ORT_WITH_TRT_EP
if (provider == OrtProvider::TensorRT) {
try {
std::filesystem::create_directories(trt_cfg.cache_dir);
// Use V2 API: key-value string map supports all options including
// dynamic batch profiles (missing from the legacy V1 struct).
std::unordered_map<std::string, std::string> kv = {
{"device_id", "0"},
{"trt_max_workspace_size", "2147483648"},
@@ -93,7 +90,7 @@ inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
trt_v2.Update(kv);
opts.AppendExecutionProvider_TensorRT_V2(*trt_v2);
std::cerr << "[" << label << "] TensorRT"
std::cerr << "[" << label << "] TensorRT EP"
<< (trt_cfg.fp16 ? " FP16" : "")
<< (trt_cfg.int8 ? " INT8" : "")
<< " cache=" << trt_cfg.cache_dir
@@ -104,11 +101,12 @@ inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
<< "\n";
return OrtProvider::TensorRT;
} catch (const Ort::Exception& e) {
std::cerr << "[" << label << "] TensorRT unavailable ("
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{};
@@ -137,10 +135,3 @@ inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
std::cerr << "[" << label << "] CPU provider\n";
return OrtProvider::CPU;
}
// Convenience: detect + apply in one call.
inline OrtProvider setup_ort_session(Ort::SessionOptions& opts,
const char* label,
const TrtConfig& trt_cfg = {}) {
return apply_ort_provider(opts, detect_ort_provider(), label, trt_cfg);
}
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);
}
src/config.hpp
+4 -4
View File
@@ -1,5 +1,5 @@
#pragma once
#include "ort_provider.hpp"
#include "inference/backend_config.hpp"
#include <string>
inline const std::string kDefaultDetectorModel = std::string(SAE_MODELS_DIR) + "/scrfd_500m_bnkps.onnx";
@@ -58,10 +58,10 @@ struct Config {
double extinction_sec{5.0}; // keep actor active this many seconds after last detection
double anneal_sec{2.0}; // merge actor windows separated by less than this into one epoch
// ── TensorRT ──────────────────────────────────────────────────────────────
// Only active when OrtProvider::TensorRT is detected.
// ── Inference backend tuning ────────────────────────────────────────────
// Consumed by the compiled-in inference backend (ORT or TRT).
// INT8 is unsafe for ArcFace without a calibration table.
TrtConfig trt{}; // fp16=true, int8=false, cache_dir="./trt_cache"
BackendConfig trt{}; // fp16=true, int8=false, cache_dir="./trt_cache"
// ── Debug output (only used when SAE_DEBUG is defined) ───────────────────
#ifdef SAE_DEBUG
src/embed_faces.cpp
+17 -28
View File
@@ -28,13 +28,10 @@
// This binary is intentionally a thin wrapper around the same ONNX models
// used by scene_analyze, so embeddings are guaranteed compatible.
#include "arcface_embedder.hpp"
#include "trt_arcface_embedder.hpp"
#include "trt_scrfd_decoder.hpp"
#include "face_utils.hpp"
#include "ort_provider.hpp"
#include "scrfd_decoder.hpp"
#include "config.hpp"
#include "face_utils.hpp"
#include "inference/face_detector.hpp"
#include "inference/face_embedder.hpp"
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
@@ -203,30 +200,22 @@ int main(int argc, char** argv) {
return 1;
}
const OrtProvider provider = detect_ort_provider();
std::cerr << "[embed_faces] inference provider: " << provider_name(provider) << "\n";
Config cfg;
cfg.detector_model = detector_model;
cfg.detector_engine = detector_engine;
cfg.arcface_model = arcface_model;
cfg.arcface_engine = arcface_engine;
cfg.detector_conf = conf;
cfg.detector_nms = nms;
std::unique_ptr<SCRFDDecoder> ort_det;
std::unique_ptr<TrtScrfdDecoder> trt_det;
std::function<std::vector<DetectedFace>(const cv::Mat&)> detect;
if (!detector_engine.empty()) {
trt_det = std::make_unique<TrtScrfdDecoder>(detector_engine, conf, nms);
detect = [&](const cv::Mat& im) { return trt_det->detect(im); };
} else {
ort_det = std::make_unique<SCRFDDecoder>(detector_model, conf, nms, provider);
detect = [&](const cv::Mat& im) { return ort_det->detect(im); };
}
// The compiled-in inference backend (ORT or TRT) is chosen by the factories.
auto detector = make_face_detector(cfg);
auto embedder = make_face_embedder(cfg);
std::unique_ptr<ArcFaceEmbedder> ort_emb;
std::unique_ptr<TrtArcFaceEmbedder> trt_emb;
std::function<Embedding(const cv::Mat&)> embed_one;
if (!arcface_engine.empty()) {
trt_emb = std::make_unique<TrtArcFaceEmbedder>(arcface_engine);
embed_one = [&](const cv::Mat& c) { return trt_emb->embed({c})[0]; };
} else {
ort_emb = std::make_unique<ArcFaceEmbedder>(arcface_model, provider);
embed_one = [&](const cv::Mat& c) { return ort_emb->embed_one(c); };
}
std::function<std::vector<DetectedFace>(const cv::Mat&)> detect =
[&](const cv::Mat& im) { return detector->detect(im); };
std::function<Embedding(const cv::Mat&)> embed_one =
[&](const cv::Mat& c) { return embedder->embed_one(c); };
// Process images and build JSON output
json output = json::array();
src/face_embedder_engine.hpp
+12 -10
View File
@@ -6,10 +6,10 @@
// of a fresh CLI invocation per image, which would reload both ONNX sessions
// every time.
#include "arcface_embedder.hpp"
#include "config.hpp"
#include "face_utils.hpp"
#include "ort_provider.hpp"
#include "scrfd_decoder.hpp"
#include "inference/face_detector.hpp"
#include "inference/face_embedder.hpp"
#include "types.hpp"
#include <opencv2/imgcodecs.hpp>
@@ -37,11 +37,13 @@ public:
float conf = 0.5f, float nms = 0.4f, int max_side = 500)
: max_side_(max_side)
{
const OrtProvider provider = detect_ort_provider();
std::cerr << "[FaceEmbedderEngine] inference provider: "
<< provider_name(provider) << "\n";
detector_ = std::make_unique<SCRFDDecoder>(detector_model, conf, nms, provider);
embedder_ = std::make_unique<ArcFaceEmbedder>(arcface_model, provider);
Config cfg;
cfg.detector_model = detector_model;
cfg.arcface_model = arcface_model;
cfg.detector_conf = conf;
cfg.detector_nms = nms;
detector_ = make_face_detector(cfg);
embedder_ = make_face_embedder(cfg);
}
FaceEmbedResult embed_path(const std::string& path) const {
@@ -104,7 +106,7 @@ public:
}
private:
std::unique_ptr<SCRFDDecoder> detector_;
std::unique_ptr<ArcFaceEmbedder> embedder_;
std::unique_ptr<IFaceDetector> detector_;
std::unique_ptr<IFaceEmbedder> embedder_;
int max_side_;
};
src/ffmpeg_decoder.hpp
+116 -41
View File
@@ -4,8 +4,10 @@ extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/hwcontext.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
}
@@ -14,13 +16,20 @@ extern "C" {
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
// ── FFmpegDecoder ─────────────────────────────────────────────────────────────
// Seek-and-decode video reader backed by FFmpeg.
//
// Hardware decode priority: NVDEC (_cuvid variants) → CPU software.
// _cuvid decoders output NV12 to system memory directly — no explicit GPU
// frame transfer is needed. swscale converts NV12/YUV → BGR24 for the rest
// Hardware decode is selected at runtime via the generic hwaccel API
// (av_hwdevice_ctx_create): the decoder probes the device types the local
// build supports, in priority order CUDA (NVIDIA) → VAAPI (AMD/Intel) →
// software. This works across GPU vendors without vendor-specific decoder
// names.
//
// Hardware decoders output frames in GPU memory (e.g. AV_PIX_FMT_CUDA,
// AV_PIX_FMT_VAAPI); av_hwframe_transfer_data copies them to a system-memory
// frame (typically NV12), then swscale converts NV12/YUV → BGR24 for the rest
// of the pipeline.
//
// Non-copyable; wrap in unique_ptr if you need to move it.
@@ -40,26 +49,9 @@ struct FFmpegDecoder {
AVStream* stream = fmt_ctx_->streams[stream_idx_];
AVCodecID cid = stream->codecpar->codec_id;
// Try NVDEC first; fall back to software on any failure
if (use_hw) {
if (const AVCodec* hwc = hw_codec_for(cid)) {
codec_ctx_ = avcodec_alloc_context3(hwc);
avcodec_parameters_to_context(codec_ctx_, stream->codecpar);
codec_ctx_->thread_count = 1;
AVDictionary* opts = nullptr;
av_dict_set(&opts, "gpu", "0", 0);
if (avcodec_open2(codec_ctx_, hwc, &opts) >= 0) {
hw_active_ = true;
std::cerr << "[FFmpegDecoder] " << path
<< " codec=" << hwc->name << " (NVDEC)\n";
} else {
avcodec_free_context(&codec_ctx_);
std::cerr << "[FFmpegDecoder] NVDEC init failed, falling back to CPU\n";
}
av_dict_free(&opts);
}
}
// Try hardware backends in priority order; fall back to software.
if (use_hw)
try_open_hw(stream, cid);
if (!hw_active_) {
const AVCodec* swc = avcodec_find_decoder(cid);
@@ -93,6 +85,7 @@ struct FFmpegDecoder {
av_frame_free(&tmp_frame_);
av_packet_free(&pkt_);
avcodec_free_context(&codec_ctx_);
if (hw_device_ctx_) av_buffer_unref(&hw_device_ctx_);
avformat_close_input(&fmt_ctx_);
}
@@ -112,6 +105,10 @@ struct FFmpegDecoder {
bool hw_active() const { return hw_active_; }
const char* codec_name() const { return codec_ctx_ ? codec_ctx_->codec->name : "unknown"; }
// Human-readable backend: "CUDA", "VAAPI", … or "CPU".
const char* hw_backend() const {
return hw_active_ ? av_hwdevice_get_type_name(hw_type_) : "CPU";
}
// Decode the frame at target_sec and return it as BGR cv::Mat.
// Returns an empty Mat at EOF.
@@ -119,7 +116,7 @@ struct FFmpegDecoder {
// Smart seek: if the target is within max_forward_sec_ ahead of the last
// decoded position, decode forward (no seek, no flush). This is dramatically
// faster for sequential sampling because avcodec_flush_buffers + re-init on
// every call is the main bottleneck — especially with NVDEC.
// every call is the main bottleneck — especially with GPU decode.
cv::Mat read_at(double target_sec) {
AVStream* stream = fmt_ctx_->streams[stream_idx_];
int64_t tgt_pts = to_stream_pts(target_sec);
@@ -139,7 +136,7 @@ struct FFmpegDecoder {
}
// Decode forward until we reach or pass target_pts.
// Convert to BGR and unref the AVFrame immediately so NVDEC surfaces
// Convert to BGR and unref the AVFrame immediately so GPU surfaces
// are returned to the pool — holding them causes surface exhaustion
// at higher sample rates.
cv::Mat out;
@@ -161,7 +158,7 @@ struct FFmpegDecoder {
last_pts_ = pts;
if (pts >= tgt_pts)
out = to_bgr(frame_);
av_frame_unref(frame_); // release NVDEC surface immediately
av_frame_unref(frame_); // release GPU surface immediately
if (!out.empty()) break;
}
}
@@ -171,12 +168,15 @@ struct FFmpegDecoder {
private:
AVFormatContext* fmt_ctx_ = nullptr;
AVCodecContext* codec_ctx_ = nullptr;
AVBufferRef* hw_device_ctx_ = nullptr;
AVFrame* frame_ = nullptr;
AVFrame* tmp_frame_ = nullptr;
AVPacket* pkt_ = nullptr;
SwsContext* sws_ctx_ = nullptr;
int stream_idx_ = -1;
bool hw_active_ = false;
AVHWDeviceType hw_type_ = AV_HWDEVICE_TYPE_NONE;
AVPixelFormat hw_pix_fmt_ = AV_PIX_FMT_NONE;
int64_t last_pts_ = AV_NOPTS_VALUE;
int64_t max_forward_pts_ = AV_NOPTS_VALUE; // set after codec opens
@@ -186,26 +186,101 @@ private:
AV_TIME_BASE_Q, s->time_base);
}
static const AVCodec* hw_codec_for(AVCodecID id) {
const char* name = nullptr;
switch (id) {
case AV_CODEC_ID_H264: name = "h264_cuvid"; break;
case AV_CODEC_ID_HEVC: name = "hevc_cuvid"; break;
case AV_CODEC_ID_AV1: name = "av1_cuvid"; break;
case AV_CODEC_ID_MPEG2VIDEO: name = "mpeg2_cuvid"; break;
case AV_CODEC_ID_MPEG4: name = "mpeg4_cuvid"; break;
case AV_CODEC_ID_VC1: name = "vc1_cuvid"; break;
default: return nullptr;
// get_format callback: tell the decoder we want the hardware surface
// format negotiated for this device. The chosen format is stashed on the
// codec context's opaque pointer so this static callback can read it.
static AVPixelFormat get_hw_format(AVCodecContext* ctx,
const AVPixelFormat* fmts) {
auto want = *static_cast<const AVPixelFormat*>(ctx->opaque);
for (const AVPixelFormat* p = fmts; *p != AV_PIX_FMT_NONE; ++p)
if (*p == want) return *p;
std::cerr << "[FFmpegDecoder] hw surface format unavailable, "
"decoder will fall back to software output\n";
return fmts[0];
}
// Probe hardware device types in priority order and open the first that
// works for this codec. Detection is fully at runtime: only device types
// compiled into the local FFmpeg are returned by av_hwdevice_iterate_types,
// and av_hwdevice_ctx_create only succeeds if a usable device is present.
void try_open_hw(AVStream* stream, AVCodecID cid) {
static const AVHWDeviceType kPriority[] = {
AV_HWDEVICE_TYPE_CUDA, // NVIDIA
AV_HWDEVICE_TYPE_VAAPI, // AMD / Intel (Linux)
};
const std::vector<AVHWDeviceType> available = available_hw_types();
const AVCodec* dec = avcodec_find_decoder(cid);
if (!dec) return;
for (AVHWDeviceType type : kPriority) {
bool present = false;
for (AVHWDeviceType a : available) present |= (a == type);
if (!present) continue;
// Find the hw pixel format this decoder advertises for this device.
AVPixelFormat pix = hw_pix_fmt_for(dec, type);
if (pix == AV_PIX_FMT_NONE) continue;
AVBufferRef* dev_ctx = nullptr;
if (av_hwdevice_ctx_create(&dev_ctx, type, nullptr, nullptr, 0) < 0)
continue; // no usable device of this type on the machine
codec_ctx_ = avcodec_alloc_context3(dec);
avcodec_parameters_to_context(codec_ctx_, stream->codecpar);
codec_ctx_->thread_count = 1;
codec_ctx_->hw_device_ctx = av_buffer_ref(dev_ctx);
hw_pix_fmt_ = pix;
codec_ctx_->opaque = &hw_pix_fmt_;
codec_ctx_->get_format = get_hw_format;
if (avcodec_open2(codec_ctx_, dec, nullptr) >= 0) {
hw_active_ = true;
hw_type_ = type;
hw_device_ctx_ = dev_ctx;
std::cerr << "[FFmpegDecoder] codec=" << dec->name
<< " hwaccel=" << av_hwdevice_get_type_name(type)
<< "\n";
return;
}
// This backend failed to open; tear down and try the next.
avcodec_free_context(&codec_ctx_);
av_buffer_unref(&dev_ctx);
hw_pix_fmt_ = AV_PIX_FMT_NONE;
std::cerr << "[FFmpegDecoder] "
<< av_hwdevice_get_type_name(type)
<< " init failed, trying next backend\n";
}
return avcodec_find_decoder_by_name(name);
}
static std::vector<AVHWDeviceType> available_hw_types() {
std::vector<AVHWDeviceType> types;
AVHWDeviceType t = AV_HWDEVICE_TYPE_NONE;
while ((t = av_hwdevice_iterate_types(t)) != AV_HWDEVICE_TYPE_NONE)
types.push_back(t);
return types;
}
// Look up the hw-surface pixel format the decoder exposes for a device type
// (e.g. AV_PIX_FMT_CUDA for CUDA, AV_PIX_FMT_VAAPI for VAAPI).
static AVPixelFormat hw_pix_fmt_for(const AVCodec* dec, AVHWDeviceType type) {
for (int i = 0;; ++i) {
const AVCodecHWConfig* cfg = avcodec_get_hw_config(dec, i);
if (!cfg) break;
if ((cfg->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) &&
cfg->device_type == type)
return cfg->pix_fmt;
}
return AV_PIX_FMT_NONE;
}
cv::Mat to_bgr(AVFrame* src) {
// _cuvid decoders output NV12 to system memory.
// Generic hwaccel would output AV_PIX_FMT_CUDA and need a transfer.
// Hardware decoders hand back GPU surfaces; transfer to system memory.
AVFrame* sw = src;
if (src->format == AV_PIX_FMT_CUDA) {
tmp_frame_->format = AV_PIX_FMT_NV12;
if (src->format == hw_pix_fmt_ && hw_pix_fmt_ != AV_PIX_FMT_NONE) {
av_frame_unref(tmp_frame_);
if (av_hwframe_transfer_data(tmp_frame_, src, 0) < 0) return {};
av_frame_copy_props(tmp_frame_, src);
sw = tmp_frame_;
src/gallery/gallery_builder.cpp
+12 -10
View File
@@ -1,8 +1,8 @@
#include "gallery_builder.hpp"
#include "arcface_embedder.hpp"
#include "config.hpp"
#include "face_utils.hpp"
#include "ort_provider.hpp"
#include "scrfd_decoder.hpp"
#include "inference/face_detector.hpp"
#include "inference/face_embedder.hpp"
#include <opencv2/imgcodecs.hpp>
@@ -31,11 +31,13 @@ static std::pair<std::string, std::string> parse_dir_name(const std::string& dir
// ── Public API ────────────────────────────────────────────────────────────────
ActorGallery build_gallery(const BuildConfig& cfg) {
const OrtProvider provider = detect_ort_provider();
std::cerr << "[build_gallery] inference provider: " << provider_name(provider) << "\n";
SCRFDDecoder decoder(cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider);
ArcFaceEmbedder arcface(cfg.arcface_model, provider);
Config icfg;
icfg.detector_model = cfg.detector_model;
icfg.arcface_model = cfg.arcface_model;
icfg.detector_conf = cfg.detector_conf;
icfg.detector_nms = cfg.detector_nms;
auto decoder = make_face_detector(icfg);
auto arcface = make_face_embedder(icfg);
ActorGallery gallery;
@@ -70,7 +72,7 @@ ActorGallery build_gallery(const BuildConfig& cfg) {
}
}
auto faces = decoder.detect(img);
auto faces = decoder->detect(img);
if (faces.empty()) {
std::cerr << " [skip] no face: " << img_file.path().filename() << "\n";
@@ -93,7 +95,7 @@ ActorGallery build_gallery(const BuildConfig& cfg) {
continue;
}
Embedding emb = arcface.embed_one(crop);
Embedding emb = arcface->embed_one(crop);
actor.embeddings.push_back(emb);
actor.source_images.push_back(img_file.path().filename().string());
src/inference/backend_config.hpp
+35
View File
@@ -0,0 +1,35 @@
#pragma once
#include <string>
// ── BackendConfig ─────────────────────────────────────────────────────────────
// Backend-neutral tuning knobs for the inference backends. Carried inside Config
// (as Config::trt) and handed to the make_face_detector / make_face_embedder
// factories. The core application sets these fields without knowing which
// backend (ONNX Runtime or TensorRT) will consume them.
//
// fp16: FP16 Tensor Core kernels — safe for both SCRFD and ArcFace.
// int8: INT8 quantisation — fast but UNSAFE for ArcFace without a
// calibration table (embedding cosine space will shift, breaking
// similarity thresholds). Safe for the SCRFD detector.
// cache_dir: TRT engines are compiled once and cached here. First run is
// slow (~30–60 s per model); every subsequent run loads instantly.
//
// Shape profile (optional, set input_name + profile_{min,opt,max} to enable):
// Shape strings are trtexec-style, e.g. "1x3x112x112". When left empty the
// backend derives a sensible default from the model.
//
// ort_cache_dir: ORT optimized-model cache. On first load ORT writes a
// pre-optimized .ort file here; subsequent loads skip graph optimization.
// Empty string = disabled. Applies to all ORT providers (ROCm, CUDA, CPU).
struct BackendConfig {
bool fp16 = true;
bool int8 = false;
std::string cache_dir = "./trt_cache";
// Per-tensor optimisation profile. All four fields must be set together.
std::string input_name; // e.g. "input.1"
std::string profile_min; // e.g. "1x3x112x112"
std::string profile_opt; // e.g. "4x3x112x112"
std::string profile_max; // e.g. "8x3x112x112"
std::string ort_cache_dir = "./ort_cache";
};
src/inference/face_detector.hpp
+27
View File
@@ -0,0 +1,27 @@
#pragma once
#include "types.hpp"
#include <memory>
#include <vector>
// ── IFaceDetector ─────────────────────────────────────────────────────────────
// Backend-agnostic face detector interface. The core application detects faces
// through this interface without knowing whether the implementation is ONNX
// Runtime (SCRFD via ORT) or raw TensorRT (a pre-built SCRFD engine).
//
// The concrete implementation is selected at compile time by CMake
// (SAE_INFERENCE_BACKEND): exactly one of backends/ort_backend.cpp or
// backends/trt_backend.cpp is compiled and provides make_face_detector().
struct Config;
struct IFaceDetector {
virtual ~IFaceDetector() = default;
// Detect all faces in a BGR image. Thread-safety is backend-defined; callers
// in this project drive a detector from a single pipeline thread.
virtual std::vector<DetectedFace> detect(const cv::Mat& img) = 0;
};
// Construct the detector for the compiled-in backend. Reads cfg.detector_model,
// cfg.detector_engine, cfg.detector_conf, cfg.detector_nms and cfg.trt.
std::unique_ptr<IFaceDetector> make_face_detector(const Config& cfg);
src/inference/face_embedder.hpp
+32
View File
@@ -0,0 +1,32 @@
#pragma once
#include "types.hpp"
#include <memory>
#include <vector>
// ── IFaceEmbedder ─────────────────────────────────────────────────────────────
// Backend-agnostic ArcFace embedder interface. The core application produces
// 512-d L2-normalised embeddings through this interface without knowing whether
// the implementation is ONNX Runtime or raw TensorRT.
//
// The concrete implementation is selected at compile time by CMake
// (SAE_INFERENCE_BACKEND): exactly one of backends/ort_backend.cpp or
// backends/trt_backend.cpp is compiled and provides make_face_embedder().
struct Config;
struct IFaceEmbedder {
virtual ~IFaceEmbedder() = default;
// Embed a batch of 112×112 BGR crops → one L2-normalised 512-d embedding
// each, parallel to the input.
virtual std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) = 0;
// Largest batch the backend accepts in a single embed() call. TRT engines
// are capped by their build profile; ORT reports the configured batch size.
virtual int max_batch() const = 0;
Embedding embed_one(const cv::Mat& crop) { return embed({crop}).front(); }
};
// Construct the embedder for the compiled-in backend. Reads cfg.arcface_model,
// cfg.arcface_engine, cfg.embed_batch_size and cfg.trt.
std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg);
src/inference/similarity.hpp
+35
View File
@@ -0,0 +1,35 @@
#pragma once
#include <memory>
#include <vector>
// ── ISimilarityEngine ─────────────────────────────────────────────────────────
// Backend-agnostic gallery similarity engine for the identity matcher.
//
// The full reference gallery (n_gallery × 512 L2-normalised embeddings) is
// uploaded to the GPU once at construction and stays resident. Per frame, the
// small query matrix (n_faces × 512) is uploaded and a single SGEMM produces the
// similarity matrix S (n_gallery × n_faces, column-major) — i.e. S[g + f*n_gal]
// is cosine_similarity(gallery[g], query[f]).
//
// The GPU math backend (cuBLAS/CUDA or rocBLAS/HIP) is selected at compile time
// by CMake (SAE_GEMM_BACKEND); backends/gemm_backend.cpp provides
// make_similarity_engine(). The core matcher node sees only this interface and
// holds no CUDA/HIP/BLAS headers.
struct ISimilarityEngine {
virtual ~ISimilarityEngine() = default;
// Largest n_faces accepted by compute() per call (bounds GPU buffer sizes).
virtual int max_faces() const = 0;
// Compute similarities for n_faces query embeddings.
// query_row_major: n_faces × 512, row fi at query + fi*512.
// Returns a pointer to host memory holding S column-major: the gallery
// similarities for face fi start at result + fi*n_gallery. The pointer is
// owned by the engine and valid until the next compute() call.
virtual const float* compute(const float* query_row_major, int n_faces) = 0;
};
// gallery_row_major: n_gallery × 512, embedding i at gallery + i*512.
std::unique_ptr<ISimilarityEngine> make_similarity_engine(
const float* gallery_row_major, int n_gallery, int max_faces);
src/main.cpp
+2 -6
View File
@@ -28,7 +28,6 @@
// --crop-context <f> bbox expansion factor for context crops (default: 1.5)
#include "config.hpp"
#include "ort_provider.hpp"
#include "types.hpp"
#include "gallery/gallery_store.hpp"
#include "nodes/frame_source_node.hpp"
@@ -139,13 +138,10 @@ int main(int argc, char** argv) {
std::atomic<bool> done{false};
const OrtProvider provider = detect_ort_provider();
std::cerr << "[main] inference provider: " << provider_name(provider) << "\n";
FrameSourceFunc source_fn {cfg};
FaceDetectorFunc detector_fn{cfg, provider};
FaceDetectorFunc detector_fn{cfg};
FaceAlignerFunc aligner_fn;
EmbedderFunc embedder_fn{cfg, provider};
EmbedderFunc embedder_fn{cfg};
FaceTrackerFunc ftracker_fn{cfg};
IdentityMatcherFunc matcher_fn {gallery, cfg};
SceneTrackerFunc tracker_fn {cfg};
src/nodes/embedder_node.hpp
+17 -25
View File
@@ -1,9 +1,8 @@
#pragma once
#include "arcface_embedder.hpp"
#include "config.hpp"
#include "ort_provider.hpp"
#include "trt_arcface_embedder.hpp"
#include "inference/face_embedder.hpp"
#include <algorithm>
#include <memory>
#include <stdexcept>
#include <string>
@@ -12,31 +11,25 @@
// KPN node: runs ArcFace on every 112×112 crop in an AlignedSceneFrame,
// producing one L2-normalised 512-dim embedding per face.
//
// Backend selection:
// --arcface-engine <path> → TrtArcFaceEmbedder (raw TensorRT, no ORT)
// otherwise → ArcFaceEmbedder (ONNX Runtime, picks best EP)
// The inference backend (ONNX Runtime or raw TensorRT) is selected at compile
// time; this node talks only to IFaceEmbedder via make_face_embedder(cfg).
//
// All crops in one frame are batched into a single forward pass (capped at
// embed_batch_size). The backends serialise themselves; we only call them
// from the single embedder thread.
// embed_batch_size). The backend serialises itself; we only call it from the
// single embedder thread.
struct EmbedderFunc {
static constexpr std::string_view label() { return "embedder"; }
explicit EmbedderFunc(const Config& cfg, OrtProvider provider)
: batch_size_(std::max(1, cfg.embed_batch_size))
explicit EmbedderFunc(const Config& cfg)
: embedder_(make_face_embedder(cfg))
, batch_size_(std::max(1, cfg.embed_batch_size))
{
if (!cfg.arcface_engine.empty()) {
trt_ = std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
if (trt_->max_batch() < static_cast<int>(batch_size_))
throw std::runtime_error(
"embed_batch_size " + std::to_string(batch_size_) +
" exceeds engine max_batch " + std::to_string(trt_->max_batch()) +
" — rebuild engine with EMBED_BATCH=" + std::to_string(batch_size_));
} else {
ort_ = std::make_unique<ArcFaceEmbedder>(
cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
}
if (embedder_->max_batch() < static_cast<int>(batch_size_))
throw std::runtime_error(
"embed_batch_size " + std::to_string(batch_size_) +
" exceeds backend max_batch " + std::to_string(embedder_->max_batch()) +
" — rebuild the engine with EMBED_BATCH=" + std::to_string(batch_size_));
}
EmbeddedSceneFrame operator()(AlignedSceneFrame af) {
@@ -49,7 +42,7 @@ struct EmbedderFunc {
for (size_t i = 0; i < crops.size(); i += batch_size_) {
const size_t end = std::min(i + batch_size_, crops.size());
std::vector<cv::Mat> chunk_crops(crops.begin() + i, crops.begin() + end);
auto chunk = trt_ ? trt_->embed(chunk_crops) : ort_->embed(chunk_crops);
auto chunk = embedder_->embed(chunk_crops);
embeddings.insert(embeddings.end(), chunk.begin(), chunk.end());
}
@@ -60,7 +53,6 @@ struct EmbedderFunc {
}
private:
std::unique_ptr<ArcFaceEmbedder> ort_;
std::unique_ptr<TrtArcFaceEmbedder> trt_;
size_t batch_size_;
std::unique_ptr<IFaceEmbedder> embedder_;
size_t batch_size_;
};
src/nodes/face_detector_node.hpp
+12 -22
View File
@@ -1,39 +1,30 @@
#pragma once
#include "scrfd_decoder.hpp"
#include "trt_scrfd_decoder.hpp"
#include "config.hpp"
#include "ort_provider.hpp"
#include "inference/face_detector.hpp"
#include <algorithm>
#include <memory>
#include <string>
// ── FaceDetectorFunc ──────────────────────────────────────────────────────────
// KPN node: runs SCRFD-500MF to detect ALL faces in a frame.
//
// Backend selection:
// --detector-engine <path> → TrtScrfdDecoder (raw TensorRT, no ORT)
// otherwise → SCRFDDecoder (ONNX Runtime)
// The inference backend (ONNX Runtime or raw TensorRT) is selected at compile
// time; this node talks only to IFaceDetector via make_face_detector(cfg).
struct FaceDetectorFunc {
static constexpr std::string_view label() { return "face_detector"; }
explicit FaceDetectorFunc(const Config& cfg, OrtProvider provider)
: max_faces_(cfg.max_faces)
explicit FaceDetectorFunc(const Config& cfg)
: detector_(make_face_detector(cfg))
, max_faces_(cfg.max_faces)
, min_face_px_(cfg.min_face_px)
{
if (!cfg.detector_engine.empty()) {
trt_ = std::make_unique<TrtScrfdDecoder>(
cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
} else {
ort_ = std::make_unique<SCRFDDecoder>(
cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
}
}
{}
SceneFrame operator()(Frame f) {
if (f.eof) return {std::move(f), {}};
auto faces = trt_ ? trt_->detect(f.image) : ort_->detect(f.image);
auto faces = detector_->detect(f.image);
// Drop faces below minimum pixel size (too small for reliable ArcFace alignment)
faces.erase(
@@ -54,8 +45,7 @@ struct FaceDetectorFunc {
}
private:
std::unique_ptr<SCRFDDecoder> ort_;
std::unique_ptr<TrtScrfdDecoder> trt_;
int max_faces_{10};
float min_face_px_{40.f};
std::unique_ptr<IFaceDetector> detector_;
int max_faces_{10};
float min_face_px_{40.f};
};
src/nodes/frame_source_node.hpp
+3 -2
View File
@@ -14,7 +14,8 @@
// ── FrameSourceFunc ───────────────────────────────────────────────────────────
// KPN source node: reads a movie file and emits one Frame per sample interval.
//
// Decode backend: FFmpeg with NVDEC (_cuvid) when available, CPU otherwise.
// Decode backend: FFmpeg hwaccel (CUDA/VAAPI, runtime-detected) when
// available, CPU otherwise.
//
// Sampling strategy: seek to the next target timestamp rather than decoding
// every frame, which is fast even for 1-FPS sampling of a 2-hour film.
@@ -39,7 +40,7 @@ struct FrameSourceFunc {
- cfg.start_sec;
int n_frames = static_cast<int>(span_s * cfg.sample_fps);
std::cerr << "[frame_source] decoder=" << decoder_->codec_name()
<< " (" << (decoder_->hw_active() ? "NVDEC" : "CPU") << ")"
<< " (" << decoder_->hw_backend() << ")"
<< " video_fps=" << decoder_->fps()
<< " start=" << cfg.start_sec << "s"
<< (end_sec_ > 0 ? " end=" + std::to_string(end_sec_) + "s" : "")
src/nodes/identity_matcher_node.hpp
+10 -88
View File
@@ -1,16 +1,15 @@
#pragma once
#include "types.hpp"
#include "config.hpp"
#include "inference/similarity.hpp"
#include "gallery/gallery_store.hpp"
#include "gallery/gallery_calibration.hpp"
#include <cublas_v2.h>
#include <cuda_runtime_api.h>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>
@@ -35,25 +34,9 @@
// Gallery scan: the full reference set (tens of thousands of 512-dim
// embeddings) is uploaded to the GPU once at construction time and stays
// resident there. Per frame, only the small query matrix (n_faces x 512) is
// uploaded and a single cublasSgemm computes the full similarity matrix
// (n_faces x N_gallery) in well under a millisecond — far faster than any
// CPU GEMM or scalar scan, making gallery-side pruning unnecessary.
namespace identity_matcher_detail {
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));
}
inline void check_cublas(cublasStatus_t s, const char* what) {
if (s != CUBLAS_STATUS_SUCCESS)
throw CudaError(std::string(what) + ": cublas error " + std::to_string(s));
}
} // namespace identity_matcher_detail
// uploaded and a single SGEMM computes the full similarity matrix in well under
// a millisecond. The GPU math backend (cuBLAS or rocBLAS) lives behind
// ISimilarityEngine (backends/gemm_backend.cpp) and is selected at compile time.
struct IdentityMatcherFunc {
static constexpr std::string_view label() { return "identity_matcher"; }
@@ -69,8 +52,6 @@ struct IdentityMatcherFunc {
, ratio_(cfg.match_ratio)
, ratio_ceil_(cfg.match_ratio_ceil)
{
using namespace identity_matcher_detail;
std::cerr << "[identity_matcher] flattening gallery embeddings...\n";
for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
for (const auto& emb : gallery_.actors[ai].embeddings) {
@@ -100,47 +81,15 @@ struct IdentityMatcherFunc {
<< gallery_.actors.size() << " actors, "
<< flat_emb_.size() << " reference embeddings\n";
// Flatten gallery into a contiguous (N x 512) row-major host buffer,
// then upload once. Row-major NxD == column-major DxN, which is the
// layout cublasSgemm wants for the transposed operand below.
std::vector<float> host_gallery(static_cast<size_t>(n_gallery_) * 512);
for (int i = 0; i < n_gallery_; ++i)
std::memcpy(host_gallery.data() + static_cast<size_t>(i) * 512,
flat_emb_[i].data(), 512 * sizeof(float));
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_gallery_), host_gallery.size() * sizeof(float)),
"cudaMalloc gallery");
check_cuda(cudaMemcpy(d_gallery_, host_gallery.data(),
host_gallery.size() * sizeof(float),
cudaMemcpyHostToDevice),
"cudaMemcpy gallery H2D");
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_query_), static_cast<size_t>(kMaxFaces) * 512 * sizeof(float)),
"cudaMalloc query");
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_sims_), static_cast<size_t>(kMaxFaces) * n_gallery_ * sizeof(float)),
"cudaMalloc sims");
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
check_cublas(cublasCreate(&handle_), "cublasCreate");
check_cublas(cublasSetStream(handle_, stream_), "cublasSetStream");
host_sims_.resize(static_cast<size_t>(kMaxFaces) * n_gallery_);
std::cerr << "[identity_matcher] gallery resident on GPU ("
<< (host_gallery.size() * sizeof(float)) / (1024 * 1024) << " MiB)\n";
}
~IdentityMatcherFunc() {
if (d_gallery_) cudaFree(d_gallery_);
if (d_query_) cudaFree(d_query_);
if (d_sims_) cudaFree(d_sims_);
if (handle_) cublasDestroy(handle_);
if (stream_) cudaStreamDestroy(stream_);
sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces);
}
MatchedSceneFrame operator()(TrackedSceneFrame tf) {
using namespace identity_matcher_detail;
if (tf.source.eof) return {std::move(tf.source), {}};
const int n_faces = static_cast<int>(tf.embeddings.size());
@@ -151,37 +100,18 @@ struct IdentityMatcherFunc {
if (n_faces > kMaxFaces)
throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");
// Build the (n_faces x 512) query matrix (row-major == col-major 512 x n_faces).
std::vector<float> host_query(static_cast<size_t>(n_faces) * 512);
for (int fi = 0; fi < n_faces; ++fi) {
std::memcpy(host_query.data() + static_cast<size_t>(fi) * 512,
tf.embeddings[fi].data(), 512 * sizeof(float));
}
check_cuda(cudaMemcpyAsync(d_query_, host_query.data(),
host_query.size() * sizeof(float),
cudaMemcpyHostToDevice, stream_),
"cudaMemcpy query H2D");
// S (N_gallery x n_faces) col-major = G(512 x N_gallery)^T * Q(512 x n_faces)
// i.e. S[g + f*N_gallery] = cosine_similarity(gallery[g], query[f]).
const float alpha = 1.f, beta = 0.f;
check_cublas(cublasSgemm(handle_, CUBLAS_OP_T, CUBLAS_OP_N,
n_gallery_, n_faces, 512,
&alpha, d_gallery_, 512, d_query_, 512,
&beta, d_sims_, n_gallery_),
"cublasSgemm");
check_cuda(cudaMemcpyAsync(host_sims_.data(), d_sims_,
static_cast<size_t>(n_gallery_) * n_faces * sizeof(float),
cudaMemcpyDeviceToHost, stream_),
"cudaMemcpy sims D2H");
check_cuda(cudaStreamSynchronize(stream_), "cudaStreamSynchronize");
// S (N_gallery × n_faces) col-major: face fi's gallery sims at sims + fi*n_gallery.
const float* host_sims = sim_engine_->compute(host_query.data(), n_faces);
for (int fi = 0; fi < n_faces; ++fi) {
const float* sims = host_sims_.data() + static_cast<size_t>(fi) * n_gallery_;
const float* sims = host_sims + static_cast<size_t>(fi) * n_gallery_;
// Per-actor best cosine similarity (max over that actor's reference embeddings)
std::vector<float> best_sim(gallery_.actors.size(),
-std::numeric_limits<float>::max());
for (int ei = 0; ei < n_gallery_; ++ei) {
@@ -190,7 +120,6 @@ struct IdentityMatcherFunc {
if (sim > best_sim[ai]) best_sim[ai] = sim;
}
// Find best and second-best actor by similarity
int best_actor = -1;
int second_actor = -1;
float best_s = -std::numeric_limits<float>::max();
@@ -240,7 +169,6 @@ struct IdentityMatcherFunc {
? cal_.probability(best_s, log_prior_odds_)
: best_s;
}
// actor_idx == -1, name == "" → unknown face
actors.push_back(std::move(ia));
}
@@ -260,11 +188,5 @@ private:
std::vector<int> flat_actor_;
int n_gallery_{0};
float* d_gallery_{nullptr}; // (n_gallery_ x 512), resident for the lifetime of this node
float* d_query_{nullptr}; // (kMaxFaces x 512)
float* d_sims_{nullptr}; // (n_gallery_ x kMaxFaces), column-major
std::vector<float> host_sims_;
cudaStream_t stream_{nullptr};
cublasHandle_t handle_{nullptr};
std::unique_ptr<ISimilarityEngine> sim_engine_;
};
src/scene_preview.cpp
+2 -6
View File
@@ -17,7 +17,6 @@
// --preview-width <px> max display width in pixels (default: 1280)
#include "config.hpp"
#include "ort_provider.hpp"
#include "types.hpp"
#include "gallery/gallery_store.hpp"
#include "nodes/frame_source_node.hpp"
@@ -110,13 +109,10 @@ int main(int argc, char** argv) {
// ── Functors ──────────────────────────────────────────────────────────────
std::atomic<bool> done{false};
const OrtProvider provider = detect_ort_provider();
std::cerr << "[main] inference provider: " << provider_name(provider) << "\n";
FrameSourceFunc source_fn {cfg};
FaceDetectorFunc detector_fn{cfg, provider};
FaceDetectorFunc detector_fn{cfg};
FaceAlignerFunc aligner_fn;
EmbedderFunc embedder_fn{cfg, provider};
EmbedderFunc embedder_fn{cfg};
FaceTrackerFunc ftracker_fn{cfg};
IdentityMatcherFunc matcher_fn {gallery, cfg};
SceneTrackerFunc tracker_fn {cfg};
src/trt_arcface_embedder.hpp
-202
View File
@@ -1,202 +0,0 @@
#pragma once
#include "face_utils.hpp"
#include "types.hpp"
#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <iostream>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>
#include <vector>
// ── TrtArcFaceEmbedder ────────────────────────────────────────────────────────
// Pure-TensorRT ArcFace runner. Loads a serialised engine built by
// scripts/build_trt_engines.sh (or any trtexec-produced .engine matching the
// ArcFace I/O contract: input Nx3x112x112 float32, output Nx512 float32 or
// float16).
//
// Skips ONNX Runtime entirely — useful on systems where ORT was built without
// the TensorRT EP (e.g. Arch's onnxruntime-opt-cuda 1.24.x).
//
// Thread-safety: a single IExecutionContext is not safe to drive from multiple
// threads concurrently; we serialise with a mutex. The KPN embedder node is
// single-threaded anyway.
namespace trt_arcface_detail {
inline std::string trim_path(const std::string& s) { return s; }
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; }
} // namespace trt_arcface_detail
struct TrtArcFaceEmbedder {
explicit TrtArcFaceEmbedder(const std::string& engine_path) {
using namespace trt_arcface_detail;
std::ifstream f(engine_path, std::ios::binary | std::ios::ate);
if (!f) throw std::runtime_error("TrtArcFaceEmbedder: cannot open " + engine_path);
const std::streamsize sz = f.tellg();
f.seekg(0);
std::vector<char> blob(sz);
f.read(blob.data(), sz);
runtime_.reset(nvinfer1::createInferRuntime(logger()));
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), sz));
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
context_.reset(engine_->createExecutionContext());
if (!context_) throw std::runtime_error("createExecutionContext failed");
// Resolve I/O tensor names + the max batch the profile permits.
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() {
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 { return max_batch_; }
std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) const {
using namespace trt_arcface_detail;
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:
struct TrtDeleter { template<class T> void operator()(T* p) const { delete p; } };
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_;
};
src/trt_scrfd_decoder.hpp
-263
View File
@@ -1,263 +0,0 @@
#pragma once
#include "trt_arcface_embedder.hpp" // pulls in CudaError/check_cuda/TrtLogger + nvinfer1/cuda headers
#include "types.hpp"
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <algorithm>
#include <array>
#include <cstring>
#include <fstream>
#include <iostream>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>
#include <vector>
// ── TrtScrfdDecoder ───────────────────────────────────────────────────────────
// Pure-TensorRT SCRFD face detector. Loads a serialised engine built by
// scripts/build_trt_engines.sh (1x3x640x640 input pinned). Post-processing
// matches scrfd_decoder.hpp byte-for-byte — only inference is swapped.
//
// Output layout (9 tensors, InsightFace export order — preserved by trtexec):
// [0..2] score_s8 / score_s16 / score_s32 (N,1)
// [3..5] bbox_s8 / bbox_s16 / bbox_s32 (N,4)
// [6..8] kps_s8 / kps_s16 / kps_s32 (N,10)
struct TrtScrfdDecoder {
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)
{
using namespace trt_arcface_detail;
std::ifstream f(engine_path, std::ios::binary | std::ios::ate);
if (!f) throw std::runtime_error("TrtScrfdDecoder: cannot open " + engine_path);
const std::streamsize sz = f.tellg();
f.seekg(0);
std::vector<char> blob(sz);
f.read(blob.data(), sz);
runtime_.reset(nvinfer1::createInferRuntime(logger()));
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), sz));
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
context_.reset(engine_->createExecutionContext());
if (!context_) throw std::runtime_error("createExecutionContext failed");
// Enumerate I/O tensors preserving engine declaration order.
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;
// Validate input shape; the engine was built with min=opt=max=1x3x640x640.
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));
// Allocate device buffer for input.
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});
// Allocate device + host buffers for each output, sized from engine.
d_outputs_.resize(n_out, nullptr);
host_outputs_.resize(n_out);
out_elem_counts_.resize(n_out, 0);
out_last_dims_.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;
out_last_dims_[oi] = last;
}
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
std::cerr << "[TrtScrfd] loaded: " << engine_path
<< " fmc=" << fmc_
<< " outputs=" << n_out << "\n";
}
~TrtScrfdDecoder() {
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) const {
using namespace trt_arcface_detail;
// Letterbox to 640×640 — identical to SCRFDDecoder.
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");
// ── Post-process: identical to SCRFDDecoder ───────────────────────────
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:
struct TrtDeleter { template<class T> void operator()(T* p) const { delete p; } };
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_;
std::vector<int> out_last_dims_;
cudaStream_t stream_ = nullptr;
mutable std::mutex mu_;
};