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:
@@ -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