feat(engine): HDF5-native galleries with embedded calibration; TensorRT backends; scene detection
Gallery format switches from JSON to HDF5 exclusively (JSON read-only kept for back-compat): save_gallery always writes HDF5, and the fitted Platt-sigmoid calibration (a, b, valid, hash) is now embedded directly in the gallery file instead of a sidecar .calib_cache.json — identity_matcher reads it from the loaded gallery and writes back only when the embeddings actually changed (hash mismatch), skipping the O(n^2) refit otherwise. Also includes: TensorRT inference backend support (ort_backend.cpp, trt_backend.cpp), gemm_backend improvements, TransNetV2-based scene-boundary detection wired through frame_source/face_tracker/main, and CMake build target updates for the new sources. Bumps the KPN submodule to feature/persistent-pipeline-reuse (push_blocking backpressure, node_ptr/node_stats introspection, ObjectVariantNodeWrapper for stateful functors) — needed by the optimizer's sae_kpn Python bindings.
This commit is contained in:
@@ -1,9 +1,11 @@
|
||||
// ── 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.
|
||||
// Similarity engine for the identity matcher. Uploads the reference gallery
|
||||
// once and computes the per-frame similarity matrix with a single SGEMM. The
|
||||
// math backend is selected at compile time by CMake (SAE_GEMM_BACKEND):
|
||||
// cuBLAS/CUDA, rocBLAS/HIP, or a portable CPU reference (SAE_GEMM_CPU). The two
|
||||
// GPU paths are the ONLY translation units that include cublas/cuda or
|
||||
// rocblas/hip headers; the CPU path pulls in no GPU headers at all and exists so
|
||||
// the pipeline can build and be tested on a machine without a GPU (CI).
|
||||
|
||||
#include "inference/similarity.hpp"
|
||||
|
||||
@@ -13,8 +15,10 @@
|
||||
#elif defined(SAE_GEMM_ROCM)
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <rocblas/rocblas.h>
|
||||
#elif defined(SAE_GEMM_CPU)
|
||||
// no external headers — portable reference implementation below
|
||||
#else
|
||||
#error "gemm_backend.cpp requires SAE_GEMM_CUDA or SAE_GEMM_ROCM to be defined"
|
||||
#error "gemm_backend.cpp requires SAE_GEMM_CUDA, SAE_GEMM_ROCM or SAE_GEMM_CPU to be defined"
|
||||
#endif
|
||||
|
||||
#include <cstring>
|
||||
@@ -26,6 +30,64 @@
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kDim = 512;
|
||||
|
||||
#if defined(SAE_GEMM_CPU)
|
||||
|
||||
// ── CPU reference engine ──────────────────────────────────────────────────────
|
||||
// Portable, dependency-free path used for CI and as the correctness oracle for
|
||||
// the GPU backends. The gallery is L2-normalised (as are the queries), so each
|
||||
// similarity is a plain dot product. S is stored column-major to match the GPU
|
||||
// backends: the gallery similarities for face fi start at result + fi*n_gallery.
|
||||
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),
|
||||
gallery_(gallery_row_major,
|
||||
gallery_row_major + static_cast<size_t>(n_gallery) * kDim)
|
||||
{
|
||||
host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);
|
||||
std::cerr << "[similarity] CPU reference engine: gallery resident in host RAM ("
|
||||
<< (gallery_.size() * sizeof(float)) / (1024 * 1024) << " MiB)\n";
|
||||
}
|
||||
|
||||
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");
|
||||
|
||||
// S(g, f) col-major = dot(gallery[g], query[f]).
|
||||
for (int f = 0; f < n_faces; ++f) {
|
||||
const float* q = query_row_major + static_cast<size_t>(f) * kDim;
|
||||
float* out = host_sims_.data() + static_cast<size_t>(f) * n_gallery_;
|
||||
for (int g = 0; g < n_gallery_; ++g) {
|
||||
const float* row = gallery_.data() + static_cast<size_t>(g) * kDim;
|
||||
float acc = 0.f;
|
||||
for (int d = 0; d < kDim; ++d) acc += row[d] * q[d];
|
||||
out[g] = acc;
|
||||
}
|
||||
}
|
||||
return host_sims_.data();
|
||||
}
|
||||
|
||||
private:
|
||||
int n_gallery_{0};
|
||||
int max_faces_{0};
|
||||
std::vector<float> gallery_; // n_gallery × 512, row-major
|
||||
std::vector<float> host_sims_; // max_faces × n_gallery, column-major
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<ISimilarityEngine> make_similarity_engine(
|
||||
const float* gallery_row_major, int n_gallery, int max_faces) {
|
||||
return std::make_unique<SimilarityEngine>(gallery_row_major, n_gallery, max_faces);
|
||||
}
|
||||
|
||||
#else // GPU backends (CUDA / ROCM)
|
||||
|
||||
struct GpuError : std::runtime_error {
|
||||
using std::runtime_error::runtime_error;
|
||||
};
|
||||
@@ -102,8 +164,6 @@ 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)
|
||||
@@ -175,3 +235,5 @@ std::unique_ptr<ISimilarityEngine> make_similarity_engine(
|
||||
const float* gallery_row_major, int n_gallery, int max_faces) {
|
||||
return std::make_unique<SimilarityEngine>(gallery_row_major, n_gallery, max_faces);
|
||||
}
|
||||
|
||||
#endif // SAE_GEMM_CPU / GPU backends
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
|
||||
#include "inference/face_detector.hpp"
|
||||
#include "inference/face_embedder.hpp"
|
||||
#include "inference/scene_detector.hpp"
|
||||
#include "backends/ort_provider.hpp"
|
||||
#include "config.hpp"
|
||||
#include "face_utils.hpp"
|
||||
@@ -336,6 +337,123 @@ private:
|
||||
bool output_is_fp16_ = false;
|
||||
};
|
||||
|
||||
// ── TransNetV2SceneDetector ───────────────────────────────────────────────────
|
||||
// ONNX Runtime TransNetV2 shot-boundary detector.
|
||||
// Input "input" : float32 [1, 100, 27, 48, 3] RGB, channels-last, 0-255
|
||||
// Output "534" : float32 [1, 100, 1] single-frame boundary logits (used)
|
||||
// Output "535" : float32 [1, 100, 1] "many-hot" head (ignored)
|
||||
// Returns per-frame sigmoid boundary probabilities. Fixed input shape, so the
|
||||
// TRT-EP profile is pinned to the single 1×100×27×48×3 shape.
|
||||
class TransNetV2SceneDetector final : public ISceneDetector {
|
||||
public:
|
||||
TransNetV2SceneDetector(const std::string& model_path,
|
||||
OrtProvider provider, BackendConfig trt_cfg)
|
||||
{
|
||||
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 shape =
|
||||
"1x" + std::to_string(kWindow) + "x" +
|
||||
std::to_string(kFrameH) + "x" + std::to_string(kFrameW) + "x3";
|
||||
trt_cfg.profile_min = shape;
|
||||
trt_cfg.profile_opt = shape;
|
||||
trt_cfg.profile_max = shape;
|
||||
}
|
||||
}
|
||||
|
||||
apply_ort_model_cache(opts, model_path, trt_cfg);
|
||||
apply_ort_provider(opts, provider, "TransNetV2", trt_cfg);
|
||||
|
||||
session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);
|
||||
|
||||
Ort::AllocatorWithDefaultOptions alloc;
|
||||
input_name_ = session_->GetInputNameAllocated(0, alloc).get();
|
||||
|
||||
// The primary boundary head is the first graph output ("534"); the
|
||||
// second ("535") is the many-hot auxiliary head we ignore. Bind both by
|
||||
// name so ORT returns them in a known order.
|
||||
const size_t n_out = session_->GetOutputCount();
|
||||
if (n_out < 1)
|
||||
throw std::runtime_error("[TransNetV2] model has no outputs");
|
||||
for (size_t i = 0; i < n_out; ++i)
|
||||
out_name_storage_.emplace_back(
|
||||
session_->GetOutputNameAllocated(i, alloc).get());
|
||||
for (auto& s : out_name_storage_) out_name_ptrs_.push_back(s.c_str());
|
||||
|
||||
std::cerr << "[TransNetV2] loaded: " << model_path
|
||||
<< " (window=" << kWindow << ")\n";
|
||||
}
|
||||
|
||||
std::vector<float> detect_window(const std::vector<cv::Mat>& window) override {
|
||||
if (static_cast<int>(window.size()) != kWindow)
|
||||
throw std::runtime_error(
|
||||
"[TransNetV2] detect_window expects exactly " +
|
||||
std::to_string(kWindow) + " frames, got " +
|
||||
std::to_string(window.size()));
|
||||
|
||||
// Pack into NHWC-per-frame contiguous buffer [100][27][48][3], RGB 0-255.
|
||||
std::vector<float> buf(
|
||||
static_cast<size_t>(kWindow) * kFrameH * kFrameW * 3);
|
||||
size_t o = 0;
|
||||
for (int f = 0; f < kWindow; ++f) {
|
||||
const cv::Mat& m = window[f];
|
||||
// Expect 48×27 BGR CV_8UC3; guard against mis-sized input.
|
||||
const cv::Mat* src = &m;
|
||||
cv::Mat resized;
|
||||
if (m.cols != kFrameW || m.rows != kFrameH || m.type() != CV_8UC3) {
|
||||
cv::Mat tmp;
|
||||
if (m.type() != CV_8UC3) m.convertTo(tmp, CV_8UC3); else tmp = m;
|
||||
cv::resize(tmp, resized, {kFrameW, kFrameH}, 0, 0, cv::INTER_AREA);
|
||||
src = &resized;
|
||||
}
|
||||
for (int y = 0; y < kFrameH; ++y) {
|
||||
const cv::Vec3b* row = src->ptr<cv::Vec3b>(y);
|
||||
for (int x = 0; x < kFrameW; ++x) {
|
||||
const cv::Vec3b& px = row[x]; // BGR
|
||||
buf[o++] = static_cast<float>(px[2]); // R
|
||||
buf[o++] = static_cast<float>(px[1]); // G
|
||||
buf[o++] = static_cast<float>(px[0]); // B
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const std::array<int64_t, 5> in_shape = {1, kWindow, kFrameH, kFrameW, 3};
|
||||
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||||
auto in_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem, buf.data(), buf.size(), in_shape.data(), in_shape.size());
|
||||
|
||||
const char* in_name_c = input_name_.c_str();
|
||||
auto outs = session_->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
&in_name_c, &in_tensor, 1,
|
||||
out_name_ptrs_.data(), out_name_ptrs_.size());
|
||||
|
||||
// Primary head: logits [1,100,1] → sigmoid.
|
||||
const float* logits = outs[0].GetTensorData<float>();
|
||||
std::vector<float> probs(kWindow);
|
||||
for (int i = 0; i < kWindow; ++i)
|
||||
probs[i] = 1.f / (1.f + std::exp(-logits[i]));
|
||||
return probs;
|
||||
}
|
||||
|
||||
private:
|
||||
Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "transnetv2"};
|
||||
std::unique_ptr<Ort::Session> session_;
|
||||
std::string input_name_;
|
||||
std::vector<std::string> out_name_storage_;
|
||||
std::vector<const char*> out_name_ptrs_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ── Factories ─────────────────────────────────────────────────────────────────
|
||||
@@ -363,3 +481,15 @@ std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg) {
|
||||
return std::make_unique<ArcFaceEmbedder>(
|
||||
cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
|
||||
}
|
||||
|
||||
std::unique_ptr<ISceneDetector> make_scene_detector(const Config& cfg) {
|
||||
if (!cfg.scene_engine.empty())
|
||||
throw std::runtime_error(
|
||||
"scene_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 << "[scene_detector] ORT backend, provider: "
|
||||
<< provider_name(provider) << "\n";
|
||||
return std::make_unique<TransNetV2SceneDetector>(
|
||||
cfg.scene_model, provider, cfg.trt);
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
|
||||
#include "inference/face_detector.hpp"
|
||||
#include "inference/face_embedder.hpp"
|
||||
#include "inference/scene_detector.hpp"
|
||||
#include "config.hpp"
|
||||
#include "face_utils.hpp"
|
||||
#include "types.hpp"
|
||||
@@ -429,6 +430,118 @@ private:
|
||||
mutable std::mutex mu_;
|
||||
};
|
||||
|
||||
// ── TrtTransNetV2SceneDetector ────────────────────────────────────────────────
|
||||
// Pure-TensorRT TransNetV2. Engine input pinned to 1×100×27×48×3 (RGB, 0-255);
|
||||
// primary boundary head is the engine's first output tensor, sigmoided here to
|
||||
// match the ORT path byte-for-byte.
|
||||
class TrtTransNetV2SceneDetector final : public ISceneDetector {
|
||||
public:
|
||||
explicit TrtTransNetV2SceneDetector(const std::string& engine_path) {
|
||||
std::vector<char> blob = read_file(engine_path, "TrtTransNetV2");
|
||||
|
||||
runtime_.reset(nvinfer1::createInferRuntime(logger()));
|
||||
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
||||
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
|
||||
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
|
||||
context_.reset(engine_->createExecutionContext());
|
||||
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
||||
|
||||
const int n_io = engine_->getNbIOTensors();
|
||||
for (int i = 0; i < n_io; ++i) {
|
||||
const char* name = engine_->getIOTensorName(i);
|
||||
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)
|
||||
input_name_ = name;
|
||||
else if (output_name_.empty())
|
||||
output_name_ = name; // first output = primary boundary head
|
||||
}
|
||||
if (input_name_.empty() || output_name_.empty())
|
||||
throw std::runtime_error("TrtTransNetV2: engine missing input/output tensor");
|
||||
|
||||
const std::size_t in_count =
|
||||
static_cast<std::size_t>(kWindow) * kFrameH * kFrameW * 3;
|
||||
const std::size_t out_count = static_cast<std::size_t>(kWindow);
|
||||
check_cuda(cudaMalloc(&d_input_, in_count * 4), "cudaMalloc input");
|
||||
check_cuda(cudaMalloc(&d_output_, out_count * 4), "cudaMalloc output");
|
||||
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
|
||||
context_->setTensorAddress(input_name_.c_str(), d_input_);
|
||||
context_->setTensorAddress(output_name_.c_str(), d_output_);
|
||||
|
||||
std::cerr << "[TrtTransNetV2] loaded: " << engine_path
|
||||
<< " (window=" << kWindow << ")\n";
|
||||
}
|
||||
|
||||
~TrtTransNetV2SceneDetector() override {
|
||||
if (stream_) cudaStreamDestroy(stream_);
|
||||
if (d_input_) cudaFree(d_input_);
|
||||
if (d_output_) cudaFree(d_output_);
|
||||
}
|
||||
|
||||
TrtTransNetV2SceneDetector(const TrtTransNetV2SceneDetector&) = delete;
|
||||
TrtTransNetV2SceneDetector& operator=(const TrtTransNetV2SceneDetector&) = delete;
|
||||
|
||||
std::vector<float> detect_window(const std::vector<cv::Mat>& window) override {
|
||||
if (static_cast<int>(window.size()) != kWindow)
|
||||
throw std::runtime_error(
|
||||
"[TrtTransNetV2] detect_window expects exactly " +
|
||||
std::to_string(kWindow) + " frames, got " +
|
||||
std::to_string(window.size()));
|
||||
|
||||
std::vector<float> buf(
|
||||
static_cast<size_t>(kWindow) * kFrameH * kFrameW * 3);
|
||||
size_t o = 0;
|
||||
for (int f = 0; f < kWindow; ++f) {
|
||||
const cv::Mat& m = window[f];
|
||||
const cv::Mat* src = &m;
|
||||
cv::Mat resized;
|
||||
if (m.cols != kFrameW || m.rows != kFrameH || m.type() != CV_8UC3) {
|
||||
cv::Mat tmp;
|
||||
if (m.type() != CV_8UC3) m.convertTo(tmp, CV_8UC3); else tmp = m;
|
||||
cv::resize(tmp, resized, {kFrameW, kFrameH}, 0, 0, cv::INTER_AREA);
|
||||
src = &resized;
|
||||
}
|
||||
for (int y = 0; y < kFrameH; ++y) {
|
||||
const cv::Vec3b* row = src->ptr<cv::Vec3b>(y);
|
||||
for (int x = 0; x < kFrameW; ++x) {
|
||||
const cv::Vec3b& px = row[x]; // BGR
|
||||
buf[o++] = static_cast<float>(px[2]); // R
|
||||
buf[o++] = static_cast<float>(px[1]); // G
|
||||
buf[o++] = static_cast<float>(px[0]); // B
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::lock_guard<std::mutex> lk(mu_);
|
||||
context_->setInputShape(input_name_.c_str(),
|
||||
nvinfer1::Dims5{1, kWindow, kFrameH, kFrameW, 3});
|
||||
|
||||
check_cuda(cudaMemcpyAsync(d_input_, buf.data(), buf.size() * 4,
|
||||
cudaMemcpyHostToDevice, stream_), "H2D input");
|
||||
if (!context_->enqueueV3(stream_))
|
||||
throw std::runtime_error("TrtTransNetV2: enqueueV3 failed");
|
||||
|
||||
std::vector<float> logits(kWindow);
|
||||
check_cuda(cudaMemcpyAsync(logits.data(), d_output_, kWindow * 4,
|
||||
cudaMemcpyDeviceToHost, stream_), "D2H output");
|
||||
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||||
|
||||
std::vector<float> probs(kWindow);
|
||||
for (int i = 0; i < kWindow; ++i)
|
||||
probs[i] = 1.f / (1.f + std::exp(-logits[i]));
|
||||
return probs;
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<nvinfer1::IRuntime, TrtDeleter> runtime_;
|
||||
std::unique_ptr<nvinfer1::ICudaEngine, TrtDeleter> engine_;
|
||||
std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
|
||||
std::string input_name_;
|
||||
std::string output_name_;
|
||||
void* d_input_ = nullptr;
|
||||
void* d_output_ = nullptr;
|
||||
cudaStream_t stream_ = nullptr;
|
||||
mutable std::mutex mu_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ── Factories ─────────────────────────────────────────────────────────────────
|
||||
@@ -453,3 +566,13 @@ std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg) {
|
||||
"-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
|
||||
return std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
|
||||
}
|
||||
|
||||
std::unique_ptr<ISceneDetector> make_scene_detector(const Config& cfg) {
|
||||
if (cfg.scene_engine.empty())
|
||||
throw std::runtime_error(
|
||||
"TRT inference backend requires a pre-built TransNetV2 engine "
|
||||
"(--scene-detector-engine / cfg.scene_engine). Build one with "
|
||||
"scripts/convert_transnetv2.py --trt, or rebuild with "
|
||||
"-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
|
||||
return std::make_unique<TrtTransNetV2SceneDetector>(cfg.scene_engine);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user