From 0ee131a692d8caca5ec485f525819be1be116cdb Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sun, 28 Jun 2026 11:50:05 +0200 Subject: [PATCH] Add AMD support via ort alternative to trt --- .gitignore | 8 +- CMakeLists.txt | 217 ++++++--- models/LVFace-B_Glint360K.onnx | 3 + scripts/make_jellyfin_gallery.py | 16 +- scripts/run_from_jellyfin.py | 9 +- src/arcface_embedder.hpp | 132 ----- src/backends/gemm_backend.cpp | 177 +++++++ .../ort_backend.cpp} | 199 ++++++-- src/{ => backends}/ort_provider.hpp | 75 ++- src/backends/trt_backend.cpp | 455 ++++++++++++++++++ src/config.hpp | 8 +- src/embed_faces.cpp | 45 +- src/face_embedder_engine.hpp | 22 +- src/ffmpeg_decoder.hpp | 157 ++++-- src/gallery/gallery_builder.cpp | 22 +- src/inference/backend_config.hpp | 35 ++ src/inference/face_detector.hpp | 27 ++ src/inference/face_embedder.hpp | 32 ++ src/inference/similarity.hpp | 35 ++ src/main.cpp | 8 +- src/nodes/embedder_node.hpp | 42 +- src/nodes/face_detector_node.hpp | 34 +- src/nodes/frame_source_node.hpp | 5 +- src/nodes/identity_matcher_node.hpp | 98 +--- src/scene_preview.cpp | 8 +- src/trt_arcface_embedder.hpp | 202 -------- src/trt_scrfd_decoder.hpp | 263 ---------- 27 files changed, 1357 insertions(+), 977 deletions(-) create mode 100644 models/LVFace-B_Glint360K.onnx delete mode 100644 src/arcface_embedder.hpp create mode 100644 src/backends/gemm_backend.cpp rename src/{scrfd_decoder.hpp => backends/ort_backend.cpp} (53%) rename src/{ => backends}/ort_provider.hpp (62%) create mode 100644 src/backends/trt_backend.cpp create mode 100644 src/inference/backend_config.hpp create mode 100644 src/inference/face_detector.hpp create mode 100644 src/inference/face_embedder.hpp create mode 100644 src/inference/similarity.hpp delete mode 100644 src/trt_arcface_embedder.hpp delete mode 100644 src/trt_scrfd_decoder.hpp diff --git a/.gitignore b/.gitignore index ee6109d..e4eb4b2 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/CMakeLists.txt b/CMakeLists.txt index e5b0ea5..480aa92 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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 , 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}" + $<$: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) diff --git a/models/LVFace-B_Glint360K.onnx b/models/LVFace-B_Glint360K.onnx new file mode 100644 index 0000000..d5af23b --- /dev/null +++ b/models/LVFace-B_Glint360K.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d834ed8e927fd35b9123b2bf97c40aad05785b1f9ecfb1c4c1f6242d38d1382 +size 455533594 diff --git a/scripts/make_jellyfin_gallery.py b/scripts/make_jellyfin_gallery.py index b382490..a559281 100644 --- a/scripts/make_jellyfin_gallery.py +++ b/scripts/make_jellyfin_gallery.py @@ -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)") diff --git a/scripts/run_from_jellyfin.py b/scripts/run_from_jellyfin.py index c06cade..c4969f4 100644 --- a/scripts/run_from_jellyfin.py +++ b/scripts/run_from_jellyfin.py @@ -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)") diff --git a/src/arcface_embedder.hpp b/src/arcface_embedder.hpp deleted file mode 100644 index c385d44..0000000 --- a/src/arcface_embedder.hpp +++ /dev/null @@ -1,132 +0,0 @@ -#pragma once -#include "ort_provider.hpp" -#include "types.hpp" -#include "face_utils.hpp" - -#include -#include -#include - -#include -#include -#include -#include - -// ── 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(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 embed(const std::vector& crops) const { - if (crops.empty()) return {}; - const int n = static_cast(crops.size()); - - // BGR→RGB, build NCHW float32 blob normalised to [-1, 1] - std::vector 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 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( - mem, reinterpret_cast(blob16.ptr()), blob16.total(), - in_shape.data(), in_shape.size()) - : Ort::Value::CreateTensor( - mem, blob.ptr(), 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 result(n); - if (output_is_fp16_) { - const auto* data16 = outs[0].GetTensorData(); - std::vector 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(); - 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 session_; - std::string input_name_; - std::string output_name_; - bool input_is_fp16_ = false; - bool output_is_fp16_ = false; -}; diff --git a/src/backends/gemm_backend.cpp b/src/backends/gemm_backend.cpp new file mode 100644 index 0000000..e172a9c --- /dev/null +++ b/src/backends/gemm_backend.cpp @@ -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 +#include +#elif defined(SAE_GEMM_ROCM) +#include +#include +#else +#error "gemm_backend.cpp requires SAE_GEMM_CUDA or SAE_GEMM_ROCM to be defined" +#endif + +#include +#include +#include +#include +#include +#include + +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(n_gallery_) * kDim; + gpu_malloc(reinterpret_cast(&d_gallery_), gallery_floats * sizeof(float)); + gpu_memcpy_h2d_sync(d_gallery_, gallery_row_major, gallery_floats * sizeof(float)); + + gpu_malloc(reinterpret_cast(&d_query_), + static_cast(max_faces_) * kDim * sizeof(float)); + gpu_malloc(reinterpret_cast(&d_sims_), + static_cast(max_faces_) * n_gallery_ * sizeof(float)); + + stream_create(&stream_); + blas_create(&handle_); + blas_set_stream(handle_, stream_); + + host_sims_.resize(static_cast(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(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(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 host_sims_; + stream_t stream_{}; + blas_handle_t handle_{}; +}; + +} // namespace + +std::unique_ptr make_similarity_engine( + const float* gallery_row_major, int n_gallery, int max_faces) { + return std::make_unique(gallery_row_major, n_gallery, max_faces); +} diff --git a/src/scrfd_decoder.hpp b/src/backends/ort_backend.cpp similarity index 53% rename from src/scrfd_decoder.hpp rename to src/backends/ort_backend.cpp index 2c6fd65..4e2c2c1 100644 --- a/src/scrfd_decoder.hpp +++ b/src/backends/ort_backend.cpp @@ -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 #include #include +#include +#include +#include +#include #include #include #include #include +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(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 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 detect(const cv::Mat& img) override { const float scale = std::min(static_cast(kInputW) / img.cols, static_cast(kInputH) / img.rows); const int new_w = static_cast(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(c * stride); const float cy = static_cast(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 out_name_storage_; std::vector 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(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 embed(const std::vector& crops) override { + if (crops.empty()) return {}; + const int n = static_cast(crops.size()); + + std::vector 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 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( + mem, reinterpret_cast(blob16.ptr()), blob16.total(), + in_shape.data(), in_shape.size()) + : Ort::Value::CreateTensor( + mem, blob.ptr(), 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 result(n); + if (output_is_fp16_) { + const auto* data16 = outs[0].GetTensorData(); + std::vector 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(); + 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 session_; + std::string input_name_; + std::string output_name_; + bool input_is_fp16_ = false; + bool output_is_fp16_ = false; +}; + +} // namespace + +// ── Factories ───────────────────────────────────────────────────────────────── + +std::unique_ptr 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( + cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt); +} + +std::unique_ptr 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( + cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size); +} diff --git a/src/ort_provider.hpp b/src/backends/ort_provider.hpp similarity index 62% rename from src/ort_provider.hpp rename to src/backends/ort_provider.hpp index f1ed1b0..7f27c26 100644 --- a/src/ort_provider.hpp +++ b/src/backends/ort_provider.hpp @@ -1,4 +1,6 @@ #pragma once +#include "inference/backend_config.hpp" + #include #include #include @@ -6,8 +8,13 @@ #include #include -// 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 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); -} diff --git a/src/backends/trt_backend.cpp b/src/backends/trt_backend.cpp new file mode 100644 index 0000000..8fb83b5 --- /dev/null +++ b/src/backends/trt_backend.cpp @@ -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 +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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 void operator()(T* p) const { delete p; } }; + +inline std::vector 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 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 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(max_batch_) * 3 * 112 * 112 * + (input_is_fp16_ ? 2 : 4); + const std::size_t out_bytes = static_cast(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 embed(const std::vector& crops) override { + if (crops.empty()) return {}; + const int n = static_cast(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 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 lk(mu_); + context_->setInputShape(input_name_.c_str(), + nvinfer1::Dims4{n, 3, 112, 112}); + + const std::size_t in_count = static_cast(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(), 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(n) * 512; + std::vector host_f32(out_count); + if (output_is_fp16_) { + std::vector 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(out_count), CV_16F, host_f16.data()); + cv::Mat dst32(1, static_cast(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 out(n); + for (int i = 0; i < n; ++i) + out[i] = l2_normalise(host_f32.data() + i * 512); + return out; + } + +private: + std::unique_ptr runtime_; + std::unique_ptr engine_; + std::unique_ptr 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 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(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(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(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 detect(const cv::Mat& img) override { + const float scale = std::min(static_cast(kInputW) / img.cols, + static_cast(kInputH) / img.rows); + const int new_w = static_cast(std::round(img.cols * scale)); + const int new_h = static_cast(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 lk(mu_); + const std::size_t in_count = static_cast(3) * kInputH * kInputW; + check_cuda(cudaMemcpyAsync(d_input_, blob.ptr(), 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 raw_boxes; + std::vector raw_scores; + std::vector> 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(c * stride); + const float cy = static_cast(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 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 keep; + cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep); + + const float img_w = static_cast(img.cols); + const float img_h = static_cast(img.rows); + + std::vector 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 runtime_; + std::unique_ptr engine_; + std::unique_ptr context_; + + float conf_threshold_; + float nms_threshold_; + int fmc_{3}; + + std::string input_name_; + std::vector output_names_; + void* d_input_ = nullptr; + std::vector d_outputs_; + mutable std::vector> host_outputs_; + std::vector out_elem_counts_; + + cudaStream_t stream_ = nullptr; + mutable std::mutex mu_; +}; + +} // namespace + +// ── Factories ───────────────────────────────────────────────────────────────── + +std::unique_ptr 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( + cfg.detector_engine, cfg.detector_conf, cfg.detector_nms); +} + +std::unique_ptr 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(cfg.arcface_engine); +} diff --git a/src/config.hpp b/src/config.hpp index efe5baa..b25fb6f 100644 --- a/src/config.hpp +++ b/src/config.hpp @@ -1,5 +1,5 @@ #pragma once -#include "ort_provider.hpp" +#include "inference/backend_config.hpp" #include 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 diff --git a/src/embed_faces.cpp b/src/embed_faces.cpp index 1a1241e..6988810 100644 --- a/src/embed_faces.cpp +++ b/src/embed_faces.cpp @@ -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 #include @@ -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 ort_det; - std::unique_ptr trt_det; - std::function(const cv::Mat&)> detect; - if (!detector_engine.empty()) { - trt_det = std::make_unique(detector_engine, conf, nms); - detect = [&](const cv::Mat& im) { return trt_det->detect(im); }; - } else { - ort_det = std::make_unique(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 ort_emb; - std::unique_ptr trt_emb; - std::function embed_one; - if (!arcface_engine.empty()) { - trt_emb = std::make_unique(arcface_engine); - embed_one = [&](const cv::Mat& c) { return trt_emb->embed({c})[0]; }; - } else { - ort_emb = std::make_unique(arcface_model, provider); - embed_one = [&](const cv::Mat& c) { return ort_emb->embed_one(c); }; - } + std::function(const cv::Mat&)> detect = + [&](const cv::Mat& im) { return detector->detect(im); }; + std::function embed_one = + [&](const cv::Mat& c) { return embedder->embed_one(c); }; // Process images and build JSON output json output = json::array(); diff --git a/src/face_embedder_engine.hpp b/src/face_embedder_engine.hpp index e89e214..a583657 100644 --- a/src/face_embedder_engine.hpp +++ b/src/face_embedder_engine.hpp @@ -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 @@ -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(detector_model, conf, nms, provider); - embedder_ = std::make_unique(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 detector_; - std::unique_ptr embedder_; + std::unique_ptr detector_; + std::unique_ptr embedder_; int max_side_; }; diff --git a/src/ffmpeg_decoder.hpp b/src/ffmpeg_decoder.hpp index 9143c04..a25a679 100644 --- a/src/ffmpeg_decoder.hpp +++ b/src/ffmpeg_decoder.hpp @@ -4,8 +4,10 @@ extern "C" { #include #include #include +#include #include #include +#include #include } @@ -14,13 +16,20 @@ extern "C" { #include #include #include +#include // ── 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(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 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 available_hw_types() { + std::vector 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_; diff --git a/src/gallery/gallery_builder.cpp b/src/gallery/gallery_builder.cpp index 1b2ff47..3543de2 100644 --- a/src/gallery/gallery_builder.cpp +++ b/src/gallery/gallery_builder.cpp @@ -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 @@ -31,11 +31,13 @@ static std::pair 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()); diff --git a/src/inference/backend_config.hpp b/src/inference/backend_config.hpp new file mode 100644 index 0000000..9756142 --- /dev/null +++ b/src/inference/backend_config.hpp @@ -0,0 +1,35 @@ +#pragma once +#include + +// ── 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"; +}; diff --git a/src/inference/face_detector.hpp b/src/inference/face_detector.hpp new file mode 100644 index 0000000..afe4786 --- /dev/null +++ b/src/inference/face_detector.hpp @@ -0,0 +1,27 @@ +#pragma once +#include "types.hpp" + +#include +#include + +// ── 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 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 make_face_detector(const Config& cfg); diff --git a/src/inference/face_embedder.hpp b/src/inference/face_embedder.hpp new file mode 100644 index 0000000..017e65f --- /dev/null +++ b/src/inference/face_embedder.hpp @@ -0,0 +1,32 @@ +#pragma once +#include "types.hpp" + +#include +#include + +// ── 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 embed(const std::vector& 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 make_face_embedder(const Config& cfg); diff --git a/src/inference/similarity.hpp b/src/inference/similarity.hpp new file mode 100644 index 0000000..c1efe5a --- /dev/null +++ b/src/inference/similarity.hpp @@ -0,0 +1,35 @@ +#pragma once +#include +#include + +// ── 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 make_similarity_engine( + const float* gallery_row_major, int n_gallery, int max_faces); diff --git a/src/main.cpp b/src/main.cpp index 4af2ec9..409731c 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -28,7 +28,6 @@ // --crop-context 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 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}; diff --git a/src/nodes/embedder_node.hpp b/src/nodes/embedder_node.hpp index 89df94c..2ef1b87 100644 --- a/src/nodes/embedder_node.hpp +++ b/src/nodes/embedder_node.hpp @@ -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 #include #include #include @@ -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 → 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(cfg.arcface_engine); - if (trt_->max_batch() < static_cast(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( - cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size); - } + if (embedder_->max_batch() < static_cast(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 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 ort_; - std::unique_ptr trt_; - size_t batch_size_; + std::unique_ptr embedder_; + size_t batch_size_; }; diff --git a/src/nodes/face_detector_node.hpp b/src/nodes/face_detector_node.hpp index fccab3f..2f5b69a 100644 --- a/src/nodes/face_detector_node.hpp +++ b/src/nodes/face_detector_node.hpp @@ -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 #include #include // ── FaceDetectorFunc ────────────────────────────────────────────────────────── // KPN node: runs SCRFD-500MF to detect ALL faces in a frame. // -// Backend selection: -// --detector-engine → 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( - cfg.detector_engine, cfg.detector_conf, cfg.detector_nms); - } else { - ort_ = std::make_unique( - 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 ort_; - std::unique_ptr trt_; - int max_faces_{10}; - float min_face_px_{40.f}; + std::unique_ptr detector_; + int max_faces_{10}; + float min_face_px_{40.f}; }; diff --git a/src/nodes/frame_source_node.hpp b/src/nodes/frame_source_node.hpp index b867159..ae66c49 100644 --- a/src/nodes/frame_source_node.hpp +++ b/src/nodes/frame_source_node.hpp @@ -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(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" : "") diff --git a/src/nodes/identity_matcher_node.hpp b/src/nodes/identity_matcher_node.hpp index 2032e86..6da9588 100644 --- a/src/nodes/identity_matcher_node.hpp +++ b/src/nodes/identity_matcher_node.hpp @@ -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 -#include - #include #include #include #include +#include #include #include @@ -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(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 host_gallery(static_cast(n_gallery_) * 512); for (int i = 0; i < n_gallery_; ++i) std::memcpy(host_gallery.data() + static_cast(i) * 512, flat_emb_[i].data(), 512 * sizeof(float)); - check_cuda(cudaMalloc(reinterpret_cast(&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(&d_query_), static_cast(kMaxFaces) * 512 * sizeof(float)), - "cudaMalloc query"); - check_cuda(cudaMalloc(reinterpret_cast(&d_sims_), static_cast(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(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(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 host_query(static_cast(n_faces) * 512); for (int fi = 0; fi < n_faces; ++fi) { std::memcpy(host_query.data() + static_cast(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(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(fi) * n_gallery_; + const float* sims = host_sims + static_cast(fi) * n_gallery_; - // Per-actor best cosine similarity (max over that actor's reference embeddings) std::vector best_sim(gallery_.actors.size(), -std::numeric_limits::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::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 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 host_sims_; - - cudaStream_t stream_{nullptr}; - cublasHandle_t handle_{nullptr}; + std::unique_ptr sim_engine_; }; diff --git a/src/scene_preview.cpp b/src/scene_preview.cpp index 8fd4236..c313570 100644 --- a/src/scene_preview.cpp +++ b/src/scene_preview.cpp @@ -17,7 +17,6 @@ // --preview-width 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 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}; diff --git a/src/trt_arcface_embedder.hpp b/src/trt_arcface_embedder.hpp deleted file mode 100644 index 4af43b1..0000000 --- a/src/trt_arcface_embedder.hpp +++ /dev/null @@ -1,202 +0,0 @@ -#pragma once -#include "face_utils.hpp" -#include "types.hpp" - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// ── 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 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(max_batch_) * 3 * 112 * 112 * - (input_is_fp16_ ? 2 : 4); - const std::size_t out_bytes = static_cast(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 embed(const std::vector& crops) const { - using namespace trt_arcface_detail; - if (crops.empty()) return {}; - const int n = static_cast(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 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 lk(mu_); - context_->setInputShape(input_name_.c_str(), - nvinfer1::Dims4{n, 3, 112, 112}); - - const std::size_t in_count = static_cast(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(), 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(n) * 512; - std::vector host_f32(out_count); - if (output_is_fp16_) { - std::vector 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(out_count), CV_16F, host_f16.data()); - cv::Mat dst32(1, static_cast(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 out(n); - for (int i = 0; i < n; ++i) - out[i] = l2_normalise(host_f32.data() + i * 512); - return out; - } - -private: - struct TrtDeleter { template void operator()(T* p) const { delete p; } }; - std::unique_ptr runtime_; - std::unique_ptr engine_; - std::unique_ptr 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_; -}; diff --git a/src/trt_scrfd_decoder.hpp b/src/trt_scrfd_decoder.hpp deleted file mode 100644 index 8488cbc..0000000 --- a/src/trt_scrfd_decoder.hpp +++ /dev/null @@ -1,263 +0,0 @@ -#pragma once -#include "trt_arcface_embedder.hpp" // pulls in CudaError/check_cuda/TrtLogger + nvinfer1/cuda headers -#include "types.hpp" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// ── 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 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(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(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(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 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(kInputW) / img.cols, - static_cast(kInputH) / img.rows); - const int new_w = static_cast(std::round(img.cols * scale)); - const int new_h = static_cast(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 lk(mu_); - const std::size_t in_count = static_cast(3) * kInputH * kInputW; - check_cuda(cudaMemcpyAsync(d_input_, blob.ptr(), 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 raw_boxes; - std::vector raw_scores; - std::vector> 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(c * stride); - const float cy = static_cast(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 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 keep; - cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep); - - const float img_w = static_cast(img.cols); - const float img_h = static_cast(img.rows); - - std::vector 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 void operator()(T* p) const { delete p; } }; - std::unique_ptr runtime_; - std::unique_ptr engine_; - std::unique_ptr context_; - - float conf_threshold_; - float nms_threshold_; - int fmc_{3}; - - std::string input_name_; - std::vector output_names_; - void* d_input_ = nullptr; - std::vector d_outputs_; - mutable std::vector> host_outputs_; - std::vector out_elem_counts_; - std::vector out_last_dims_; - - cudaStream_t stream_ = nullptr; - mutable std::mutex mu_; -};