From 41a277bc1909d51bff4b4ae12fc4c8dcab5d6da5 Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sun, 19 Jul 2026 19:04:03 +0200 Subject: [PATCH] feat(engine): HDF5-native galleries with embedded calibration; TensorRT backends; scene detection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CMakeLists.txt | 57 +++++- external/KPN | 2 +- src/backends/gemm_backend.cpp | 78 ++++++++- src/backends/ort_backend.cpp | 130 ++++++++++++++ src/backends/trt_backend.cpp | 123 +++++++++++++ src/config.hpp | 85 ++++++++- src/ffmpeg_decoder.hpp | 24 ++- src/gallery/gallery_calibration.hpp | 75 ++++---- src/gallery/gallery_store.cpp | 193 +++++++++++++++++--- src/gallery/gallery_store.hpp | 17 +- src/main.cpp | 261 +++++++++++++++++++--------- src/nodes/face_detector_node.hpp | 9 +- src/nodes/face_tracker_node.hpp | 79 ++++++++- src/nodes/frame_source_node.hpp | 69 ++++---- src/nodes/identity_matcher_node.hpp | 62 ++++++- src/nodes/preview_node.hpp | 41 +++++ src/nodes/scene_tracker_node.hpp | 4 + src/scene_preview.cpp | 34 +++- src/types.hpp | 24 ++- 19 files changed, 1151 insertions(+), 216 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 480aa92..ee53db7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,10 +35,11 @@ set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/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 +# CPU → portable reference GEMM (no GPU; CI / testing) set(SAE_INFERENCE_BACKEND "ORT" CACHE STRING "Inference backend: ORT | TRT") -set(SAE_GEMM_BACKEND "ROCM" CACHE STRING "Gallery GEMM backend: ROCM | CUDA") +set(SAE_GEMM_BACKEND "ROCM" CACHE STRING "Gallery GEMM backend: ROCM | CUDA | CPU") set_property(CACHE SAE_INFERENCE_BACKEND PROPERTY STRINGS ORT TRT) -set_property(CACHE SAE_GEMM_BACKEND PROPERTY STRINGS ROCM CUDA) +set_property(CACHE SAE_GEMM_BACKEND PROPERTY STRINGS ROCM CUDA CPU) # Enable the ORT TensorRT/CUDA execution providers inside the ORT inference # backend (only meaningful when ORT was built with the TensorRT EP). Off by @@ -63,8 +64,8 @@ endif() 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}')") +if(NOT SAE_GEMM_BACKEND MATCHES "^(ROCM|CUDA|CPU)$") + message(FATAL_ERROR "SAE_GEMM_BACKEND must be ROCM, CUDA or CPU (got '${SAE_GEMM_BACKEND}')") endif() # CUDA runtime is needed by both TRT inference and CUDA GEMM — find it once. @@ -133,7 +134,16 @@ else() # TRT endif() # ── GEMM backend dependency: builds the `gemm_backend` object lib ────────────── -if(SAE_GEMM_BACKEND STREQUAL "CUDA") +if(SAE_GEMM_BACKEND STREQUAL "CPU") + # Portable reference GEMM: no GPU libraries, no headers. Used for CI and as + # the correctness oracle for the CUDA/ROCm backends. + message(STATUS "GEMM backend: CPU (portable reference, no GPU)") + + 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_compile_definitions(gemm_backend PRIVATE SAE_GEMM_CPU) +elseif(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) @@ -223,12 +233,16 @@ set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models" # 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. +# HDF5 (C++) — gallery fast-load path + embedding dump. Found here so sae_gallery +# (gallery_store.cpp) can link it; scene_analyze/dump_embeddings reuse the same vars. +find_package(HDF5 REQUIRED COMPONENTS CXX) + 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_include_directories(sae_gallery PUBLIC src ${HDF5_INCLUDE_DIRS}) target_link_libraries(sae_gallery PUBLIC kpn ${OpenCV_LIBS} @@ -236,6 +250,7 @@ target_link_libraries(sae_gallery PUBLIC inference_backend gemm_backend ffmpeg_libs + ${HDF5_CXX_LIBRARIES} ) target_compile_definitions(sae_gallery PUBLIC SAE_MODELS_DIR="${SAE_MODELS_DIR}" @@ -249,15 +264,34 @@ target_link_libraries(embed_faces PRIVATE sae_gallery) nanobind_add_module(sae_embed src/python_bindings.cpp) target_link_libraries(sae_embed PRIVATE sae_gallery) +# ── sae_kpn — Python module: run the real downstream nodes over dumped embeddings ─ +# Assembles face_tracker/identity_matcher/scene_tracker in a Python-driven KPN +# network (KPN_BUILD_PYTHON is enabled per-TU inside the .cpp). Powers the +# threshold-sweep optimizer in scripts/optimizer/. +nanobind_add_module(sae_kpn src/kpn_bindings.cpp) +target_link_libraries(sae_kpn PRIVATE sae_gallery) + +# HDF5 already found above (before sae_gallery); vars HDF5_CXX_LIBRARIES / _INCLUDE_DIRS +# are reused by scene_analyze / dump_embeddings below. + # ── analyze — main analysis binary ─────────────────────────────────────────── add_executable(scene_analyze src/main.cpp) -target_link_libraries(scene_analyze PRIVATE sae_gallery) +target_link_libraries(scene_analyze PRIVATE sae_gallery ${HDF5_CXX_LIBRARIES}) +target_include_directories(scene_analyze PRIVATE ${HDF5_INCLUDE_DIRS}) # ── analyze_debug — same binary with debug frame/crop output ───────────────── add_executable(scene_analyze_debug src/main.cpp) -target_link_libraries(scene_analyze_debug PRIVATE sae_gallery) +target_link_libraries(scene_analyze_debug PRIVATE sae_gallery ${HDF5_CXX_LIBRARIES}) +target_include_directories(scene_analyze_debug PRIVATE ${HDF5_INCLUDE_DIRS}) target_compile_definitions(scene_analyze_debug PRIVATE SAE_DEBUG=1) +# ── dump_embeddings — standalone embedding dumper, NO gallery/matcher ───────── +# Front-half only (decode→detect→align→embed→HDF5) for the optimizer replay corpus +# and model bake-off. Skips gallery load + calibration (~24s/run faster). +add_executable(dump_embeddings src/dump_embeddings.cpp) +target_link_libraries(dump_embeddings PRIVATE sae_gallery ${HDF5_CXX_LIBRARIES}) +target_include_directories(dump_embeddings PRIVATE ${HDF5_INCLUDE_DIRS}) + # ── scene_preview — live annotated display while analysing ─────────────────── add_executable(scene_preview src/scene_preview.cpp) target_link_libraries(scene_preview PRIVATE sae_gallery) @@ -273,5 +307,12 @@ if(SAE_WEB_DEBUG) kpn_target_enable_web_debug(scene_preview) endif() +# ── Tests ───────────────────────────────────────────────────────────────────── +option(SAE_BUILD_TESTS "Build unit tests (GPU-free)" OFF) +if(SAE_BUILD_TESTS) + enable_testing() + add_subdirectory(tests) +endif() + message(STATUS "OpenCV ${OpenCV_VERSION} found") message(STATUS "Models dir: ${SAE_MODELS_DIR}") diff --git a/external/KPN b/external/KPN index 949c813..4b6e498 160000 --- a/external/KPN +++ b/external/KPN @@ -1 +1 @@ -Subproject commit 949c8134efb38030b47387044f5929e64745be39 +Subproject commit 4b6e498ba7e70a34cc0b57638f9e56a43b7f41ae diff --git a/src/backends/gemm_backend.cpp b/src/backends/gemm_backend.cpp index e172a9c..ac1278c 100644 --- a/src/backends/gemm_backend.cpp +++ b/src/backends/gemm_backend.cpp @@ -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 #include +#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 @@ -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(n_gallery) * kDim) + { + host_sims_.resize(static_cast(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(f) * kDim; + float* out = host_sims_.data() + static_cast(f) * n_gallery_; + for (int g = 0; g < n_gallery_; ++g) { + const float* row = gallery_.data() + static_cast(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 gallery_; // n_gallery × 512, row-major + std::vector host_sims_; // max_faces × n_gallery, column-major +}; + +} // 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); +} + +#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 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); } + +#endif // SAE_GEMM_CPU / GPU backends diff --git a/src/backends/ort_backend.cpp b/src/backends/ort_backend.cpp index 4e2c2c1..6f78b40 100644 --- a/src/backends/ort_backend.cpp +++ b/src/backends/ort_backend.cpp @@ -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(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 detect_window(const std::vector& window) override { + if (static_cast(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 buf( + static_cast(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(y); + for (int x = 0; x < kFrameW; ++x) { + const cv::Vec3b& px = row[x]; // BGR + buf[o++] = static_cast(px[2]); // R + buf[o++] = static_cast(px[1]); // G + buf[o++] = static_cast(px[0]); // B + } + } + } + + const std::array in_shape = {1, kWindow, kFrameH, kFrameW, 3}; + auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + auto in_tensor = Ort::Value::CreateTensor( + 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(); + std::vector 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 session_; + std::string input_name_; + std::vector out_name_storage_; + std::vector out_name_ptrs_; +}; + } // namespace // ── Factories ───────────────────────────────────────────────────────────────── @@ -363,3 +481,15 @@ std::unique_ptr make_face_embedder(const Config& cfg) { return std::make_unique( cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size); } + +std::unique_ptr 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( + cfg.scene_model, provider, cfg.trt); +} diff --git a/src/backends/trt_backend.cpp b/src/backends/trt_backend.cpp index 8fb83b5..b6e8459 100644 --- a/src/backends/trt_backend.cpp +++ b/src/backends/trt_backend.cpp @@ -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 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(kWindow) * kFrameH * kFrameW * 3; + const std::size_t out_count = static_cast(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 detect_window(const std::vector& window) override { + if (static_cast(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 buf( + static_cast(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(y); + for (int x = 0; x < kFrameW; ++x) { + const cv::Vec3b& px = row[x]; // BGR + buf[o++] = static_cast(px[2]); // R + buf[o++] = static_cast(px[1]); // G + buf[o++] = static_cast(px[0]); // B + } + } + } + + std::lock_guard 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 logits(kWindow); + check_cuda(cudaMemcpyAsync(logits.data(), d_output_, kWindow * 4, + cudaMemcpyDeviceToHost, stream_), "D2H output"); + check_cuda(cudaStreamSynchronize(stream_), "stream sync"); + + std::vector 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 runtime_; + std::unique_ptr engine_; + std::unique_ptr 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 make_face_embedder(const Config& cfg) { "-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly."); return std::make_unique(cfg.arcface_engine); } + +std::unique_ptr 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(cfg.scene_engine); +} diff --git a/src/config.hpp b/src/config.hpp index 805574a..f81ef71 100644 --- a/src/config.hpp +++ b/src/config.hpp @@ -3,7 +3,8 @@ #include inline const std::string kDefaultDetectorModel = std::string(SAE_MODELS_DIR) + "/scrfd_500m_bnkps.onnx"; -inline const std::string kDefaultArcfaceModel = std::string(SAE_MODELS_DIR) + "/arcface_w600k_r50.onnx"; +inline const std::string kDefaultArcfaceModel = std::string(SAE_MODELS_DIR) + "/LVFace-B_Glint360K.onnx"; +inline const std::string kDefaultSceneModel = std::string(SAE_MODELS_DIR) + "/transnetv2.onnx"; enum class Verbosity { minimal, // actor names + merged time windows only @@ -21,6 +22,10 @@ struct Config { std::string output_path; // annotations.json Verbosity verbosity{Verbosity::minimal}; + // When set, tee the embedder output to an HDF5 dump (schema: + // scripts/optimizer/SCHEMA.md) for offline threshold-sweep replay via sae_kpn. + std::string dump_embeddings_path; + // ── Sampling ───────────────────────────────────────────────────────────── float sample_fps{1.0f}; // frames to analyse per second of movie float max_decode_fps{0.f}; // wall-clock cap on source decode rate (0 = uncapped) @@ -40,7 +45,12 @@ struct Config { std::string arcface_engine; // optional path to a pre-built TRT engine; bypasses ORT int embed_batch_size{4}; // max faces per ORT Run() call — bounds per-call latency float match_prior{0.5f}; // base-rate prior; 0.5 = use calibrated sigmoid directly - float prob_threshold{0.99f}; // posterior P(match | sim, prior) threshold + // prob_threshold tuned by Differential Evolution against Amazon X-Ray per-scene + // presence over 4 films, per-second metric (see docs/rep4-optimizer-results.md). + // Best model+mode: LVFace-B_Glint360K, full gallery, expansion on. Supersedes the + // earlier 9-film scene-union-metric tuning (0.76) — that metric is now known to + // have hidden out-of-cast false positives (see docs/optimizer-experiments.md). + float prob_threshold{0.754f}; // posterior P(match | sim, prior) threshold float match_threshold{0.45f}; // cosine distance hard ceiling fallback (no calibration) float match_ratio{0.80f}; // ratio test fallback: accept if best/second < ratio float match_ratio_ceil{0.65f}; // ratio test only fires below this absolute distance @@ -48,15 +58,82 @@ struct Config { // ── Cut detection ──────────────────────────────────────────────────────── float cut_threshold{0.70f}; // grayscale histogram correlation below this → hard cut + // ── Scene detection (TransNetV2, opt-in) ───────────────────────────────── + // When enabled, the source decodes densely (native FPS) and a decimator + // splits the stream: full-res 1-FPS frames to the face pipeline, and a + // downscaled dense stream to the TransNetV2 scene detector. Shot boundaries + // it finds are surfaced as Frame::is_scene_boundary. This is separate from + // the always-on histogram cut, which flags intra-scene camera-angle changes. + bool scene_detect{false}; // master switch (--scene-detect) + std::string scene_model; // TransNetV2 .onnx (default set in main) + std::string scene_engine; // optional pre-built TRT .engine; bypasses ORT + float scene_threshold{0.60f}; // sigmoid boundary prob above this → boundary + // (this export's non-boundary baseline sits + // at ~0.50; real boundaries spike to ~0.7+) + int scene_stride{50}; // frames advanced between windows (≤ kWindow) + + // Dense-decode throughput knobs (only active with scene_detect). Dense decode + // of every native-rate frame is the pipeline's cost driver; these trade a + // little boundary precision for a large speedup. + // scene_decode_fps: rate the source decodes at in dense mode. Lower = + // fewer frames decoded. TransNetV2 tolerates ~12fps; boundary timestamps + // stay correct (keyed off each frame's real timestamp). 0 = native fps. + // dense_scale: downscale factor applied to decoded frames in dense mode + // (0 +#include #include #include #include @@ -35,7 +36,15 @@ extern "C" { // Non-copyable; wrap in unique_ptr if you need to move it. struct FFmpegDecoder { - explicit FFmpegDecoder(const std::string& path, bool use_hw = true) { + // out_scale in (0,1] downscales decoded frames (applied in the sws_scale + // colour conversion, so it's nearly free). 1.0 = native resolution. + explicit FFmpegDecoder(const std::string& path, bool use_hw = true, + float out_scale = 1.0f) + : out_scale_(out_scale > 0.f && out_scale <= 1.f ? out_scale : 1.0f) + { + // Quiet FFmpeg's own logging (e.g. the harmless "Could not dynamically + // load CUDA" emitted while probing hwaccels before VAAPI succeeds). + av_log_set_level(AV_LOG_ERROR); if (avformat_open_input(&fmt_ctx_, path.c_str(), nullptr, nullptr) < 0) throw std::runtime_error("[FFmpegDecoder] cannot open: " + path); if (avformat_find_stream_info(fmt_ctx_, nullptr) < 0) @@ -179,6 +188,7 @@ private: 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 + float out_scale_ = 1.0f; // decoded-frame downscale (0,1] int64_t to_stream_pts(double sec) const { AVStream* s = fmt_ctx_->streams[stream_idx_]; @@ -289,13 +299,21 @@ private: const int w = sw->width; const int h = sw->height; + // Optional downscale, folded into the colour conversion (near-free). + // Round to even dimensions for swscale/codec friendliness. + int out_w = w, out_h = h; + if (out_scale_ < 1.0f) { + out_w = std::max(2, (static_cast(w * out_scale_) / 2) * 2); + out_h = std::max(2, (static_cast(h * out_scale_) / 2) * 2); + } + sws_ctx_ = sws_getCachedContext(sws_ctx_, w, h, static_cast(sw->format), - w, h, AV_PIX_FMT_BGR24, + out_w, out_h, AV_PIX_FMT_BGR24, SWS_BILINEAR, nullptr, nullptr, nullptr); if (!sws_ctx_) return {}; - cv::Mat out(h, w, CV_8UC3); + cv::Mat out(out_h, out_w, CV_8UC3); uint8_t* dst_data[1] = { out.data }; int dst_linesize[1] = { static_cast(out.step) }; sws_scale(sws_ctx_, diff --git a/src/gallery/gallery_calibration.hpp b/src/gallery/gallery_calibration.hpp index 5784272..acf761d 100644 --- a/src/gallery/gallery_calibration.hpp +++ b/src/gallery/gallery_calibration.hpp @@ -335,63 +335,48 @@ inline uint64_t hash_gallery_embeddings( return h; } -// Calibrates the gallery, caching the fitted (a, b, valid) result on disk -// keyed by a hash of the reference embeddings. The O(n^2) pairwise fit only -// re-runs when the gallery's embeddings/actor assignments actually change. +// Calibrates the gallery, reusing (cached_a, cached_b, cached_valid) if +// cached_hash matches a fresh hash of the current embeddings/actor +// assignments — the O(n^2) pairwise fit only re-runs when they actually +// change. Distinct from calibrate_gallery_cached's old sidecar-JSON-file +// design: the cache now lives in the gallery HDF5 itself (ActorGallery:: +// calib_*, see gallery_store.hpp), so this takes the previous values +// in-memory rather than a file path. Sets `recomputed` so the caller (which +// holds the open gallery file/struct) knows whether it needs to persist the +// refreshed values back. inline GalleryCalibration calibrate_gallery_cached( const std::vector& flat_emb, const std::vector& flat_actor, - const std::string& cache_path) + float cached_a, + float cached_b, + bool cached_valid, + uint64_t cached_hash, + const std::string& curve_base_path, + bool& recomputed) { uint64_t hash = hash_gallery_embeddings(flat_emb, flat_actor); + recomputed = false; - std::string base_path = cache_path; - constexpr std::string_view kJsonExt = ".json"; - if (base_path.size() >= kJsonExt.size() && - base_path.compare(base_path.size() - kJsonExt.size(), kJsonExt.size(), kJsonExt) == 0) - base_path.resize(base_path.size() - kJsonExt.size()); - - std::ifstream in(cache_path); - if (in.is_open()) { - try { - nlohmann::json j; - in >> j; - if (j.at("hash").get() == hash) { - GalleryCalibration cal; - cal.a = j.at("a").get(); - cal.b = j.at("b").get(); - cal.valid = j.at("valid").get(); - std::cerr << "[calibration] using cached calibration from " - << cache_path << " (a=" << cal.a << " b=" << cal.b - << " valid=" << cal.valid << ")\n"; - save_calibration_curve(cal, base_path); - return cal; - } - std::cerr << "[calibration] cache at " << cache_path - << " is stale, recomputing\n"; - } catch (const std::exception&) { - std::cerr << "[calibration] cache at " << cache_path - << " is unreadable, recomputing\n"; - } + if (cached_hash != 0 && cached_hash == hash) { + GalleryCalibration cal{cached_a, cached_b, cached_valid}; + std::cerr << "[calibration] using cached calibration from gallery" + << " (a=" << cal.a << " b=" << cal.b + << " valid=" << cal.valid << ")\n"; + if (!curve_base_path.empty()) save_calibration_curve(cal, curve_base_path); + return cal; } + if (cached_hash != 0) + std::cerr << "[calibration] cached calibration is stale (embeddings changed), " + "recomputing\n"; auto t0 = std::chrono::steady_clock::now(); GalleryCalibration cal = calibrate_gallery(flat_emb, flat_actor); auto t1 = std::chrono::steady_clock::now(); - double secs = std::chrono::duration(t1 - t0).count(); - std::cerr << "[calibration] fit took " << secs << "s for " + std::cerr << "[calibration] fit took " + << std::chrono::duration(t1 - t0).count() << "s for " << flat_emb.size() << " embeddings\n"; - nlohmann::json j; - j["hash"] = hash; - j["a"] = cal.a; - j["b"] = cal.b; - j["valid"] = cal.valid; - j["fit_secs"] = secs; - std::ofstream out(cache_path); - if (out.is_open()) out << j.dump(2) << "\n"; - - save_calibration_curve(cal, base_path); - + if (!curve_base_path.empty()) save_calibration_curve(cal, curve_base_path); + recomputed = true; return cal; } diff --git a/src/gallery/gallery_store.cpp b/src/gallery/gallery_store.cpp index 1e688a1..798c351 100644 --- a/src/gallery/gallery_store.cpp +++ b/src/gallery/gallery_store.cpp @@ -1,6 +1,7 @@ #include "gallery_store.hpp" #include +#include #include #include #include @@ -8,7 +9,168 @@ using json = nlohmann::json; +// ── HDF5 (see gallery_store.hpp for the full layout) ───────────────────────── +// A 170MB gallery JSON parses in ~18s (nlohmann). The same data as HDF5 loads in +// ~1s — a big win for the optimizer, which reloads the gallery per replay subprocess. +static bool ends_with(const std::string& s, const std::string& suf) { + return s.size() >= suf.size() && + s.compare(s.size() - suf.size(), suf.size(), suf) == 0; +} + +static std::vector read_str_dataset(H5::H5File& file, const char* name, hsize_t a) { + H5::DataSet ds = file.openDataSet(name); + H5::StrType st = ds.getStrType(); + std::vector out(a); + if (st.isVariableStr()) { + std::vector raw(a); + ds.read(raw.data(), st); + for (hsize_t i = 0; i < a; ++i) { out[i] = raw[i] ? raw[i] : ""; } + H5::DataSpace sp = ds.getSpace(); + H5Dvlen_reclaim(st.getId(), sp.getId(), H5P_DEFAULT, raw.data()); + } + return out; +} + +static ActorGallery load_gallery_hdf5(const std::string& path) { + std::cerr << "[gallery] loading " << path << " (HDF5)..." << std::flush; + auto t0 = std::chrono::steady_clock::now(); + + H5::H5File file(path, H5F_ACC_RDONLY); + H5::DataSet emb_ds = file.openDataSet("embeddings"); + hsize_t dims[2]; + emb_ds.getSpace().getSimpleExtentDims(dims); // [N, 512] + const hsize_t N = dims[0]; + if (dims[1] != 512) throw std::runtime_error("gallery HDF5: embedding dim != 512"); + + std::vector flat(N * 512); + emb_ds.read(flat.data(), H5::PredType::NATIVE_FLOAT); + + H5::DataSet off_ds = file.openDataSet("offset"); + hsize_t adim[1]; + off_ds.getSpace().getSimpleExtentDims(adim); + const hsize_t A = adim[0]; + std::vector offset(A); + off_ds.read(offset.data(), H5::PredType::NATIVE_INT64); + std::vector count(A); + file.openDataSet("count").read(count.data(), H5::PredType::NATIVE_INT32); + + auto imdb = read_str_dataset(file, "imdb_id", A); + auto tmdb = read_str_dataset(file, "tmdb_id", A); + auto jf = read_str_dataset(file, "jellyfin_id", A); + auto name = read_str_dataset(file, "name", A); + std::vector src_images; + if (file.nameExists("source_images")) + src_images = read_str_dataset(file, "source_images", N); + + ActorGallery gallery; + gallery.actors.reserve(A); + for (hsize_t a = 0; a < A; ++a) { + ActorGallery::Actor actor; + actor.imdb_id = imdb[a]; actor.tmdb_id = tmdb[a]; + actor.jellyfin_id = jf[a]; actor.name = name[a]; + for (int32_t e = 0; e < count[a]; ++e) { + hsize_t row = offset[a] + e; + Embedding emb; + std::copy_n(flat.data() + row * 512, 512, emb.begin()); + actor.embeddings.push_back(emb); + if (!src_images.empty()) + actor.source_images.push_back(src_images[row]); + } + gallery.actors.push_back(std::move(actor)); + } + + if (file.nameExists("calibration")) { + H5::Group cal = file.openGroup("calibration"); + cal.openAttribute("a").read(H5::PredType::NATIVE_FLOAT, &gallery.calib_a); + cal.openAttribute("b").read(H5::PredType::NATIVE_FLOAT, &gallery.calib_b); + int8_t valid = 0; + cal.openAttribute("valid").read(H5::PredType::NATIVE_INT8, &valid); + gallery.calib_valid = valid != 0; + cal.openAttribute("hash").read(H5::PredType::NATIVE_UINT64, &gallery.calib_hash); + } + + auto t1 = std::chrono::steady_clock::now(); + std::cerr << " built " << A << " actors / " << N << " embeddings in " + << std::chrono::duration(t1 - t0).count() << "s"; + if (gallery.calib_hash != 0) + std::cerr << " (calibration cached: a=" << gallery.calib_a + << " b=" << gallery.calib_b << " valid=" << gallery.calib_valid << ")"; + std::cerr << "\n"; + return gallery; +} + +static void write_str_dataset(H5::H5File& file, const char* name, + const std::vector& values) { + H5::StrType str_t(H5::PredType::C_S1, H5T_VARIABLE); + hsize_t n = values.size(); + H5::DataSpace space(1, &n); + H5::DataSet ds = file.createDataSet(name, str_t, space); + std::vector raw(n); + for (hsize_t i = 0; i < n; ++i) raw[i] = values[i].c_str(); + ds.write(raw.data(), str_t); +} + +static void save_gallery_hdf5(const std::string& path, const ActorGallery& gallery) { + H5::H5File file(path, H5F_ACC_TRUNC); + + std::vector flat; + std::vector offset; + std::vector count; + std::vector imdb, tmdb, jf, name, src_images; + int64_t row = 0; + for (const auto& a : gallery.actors) { + offset.push_back(row); + count.push_back(static_cast(a.embeddings.size())); + row += static_cast(a.embeddings.size()); + for (size_t i = 0; i < a.embeddings.size(); ++i) { + flat.insert(flat.end(), a.embeddings[i].begin(), a.embeddings[i].end()); + src_images.push_back(i < a.source_images.size() ? a.source_images[i] : ""); + } + imdb.push_back(a.imdb_id); tmdb.push_back(a.tmdb_id); + jf.push_back(a.jellyfin_id); name.push_back(a.name); + } + + hsize_t N = flat.size() / 512; + hsize_t emb_dims[2] = {N, 512}; + H5::DataSpace emb_space(2, emb_dims); + file.createDataSet("embeddings", H5::PredType::NATIVE_FLOAT, emb_space) + .write(flat.data(), H5::PredType::NATIVE_FLOAT); + + hsize_t A = gallery.actors.size(); + H5::DataSpace a_space(1, &A); + file.createDataSet("offset", H5::PredType::NATIVE_INT64, a_space) + .write(offset.data(), H5::PredType::NATIVE_INT64); + file.createDataSet("count", H5::PredType::NATIVE_INT32, a_space) + .write(count.data(), H5::PredType::NATIVE_INT32); + + write_str_dataset(file, "imdb_id", imdb); + write_str_dataset(file, "tmdb_id", tmdb); + write_str_dataset(file, "jellyfin_id", jf); + write_str_dataset(file, "name", name); + write_str_dataset(file, "source_images", src_images); + + if (gallery.calib_hash != 0) { + H5::Group cal = file.createGroup("calibration"); + H5::DataSpace scalar(H5S_SCALAR); + cal.createAttribute("a", H5::PredType::NATIVE_FLOAT, scalar) + .write(H5::PredType::NATIVE_FLOAT, &gallery.calib_a); + cal.createAttribute("b", H5::PredType::NATIVE_FLOAT, scalar) + .write(H5::PredType::NATIVE_FLOAT, &gallery.calib_b); + int8_t valid = gallery.calib_valid ? 1 : 0; + cal.createAttribute("valid", H5::PredType::NATIVE_INT8, scalar) + .write(H5::PredType::NATIVE_INT8, &valid); + cal.createAttribute("hash", H5::PredType::NATIVE_UINT64, scalar) + .write(H5::PredType::NATIVE_UINT64, &gallery.calib_hash); + } + + std::cerr << "[gallery] saved " << A << " actors / " << N + << " embeddings to " << path << " (HDF5)\n"; +} + ActorGallery load_gallery(const std::string& path) { + if (ends_with(path, ".h5") || ends_with(path, ".hdf5")) + return load_gallery_hdf5(path); + std::ifstream f(path); if (!f.is_open()) throw std::runtime_error("load_gallery: cannot open " + path); @@ -53,28 +215,15 @@ ActorGallery load_gallery(const std::string& path) { return gallery; } +// Always writes HDF5. If `path` doesn't already end in .h5/.hdf5, the +// extension is replaced (galleries are never written as JSON anymore). void save_gallery(const std::string& path, const ActorGallery& gallery) { - json j; - j["actors"] = json::array(); - - for (const auto& actor : gallery.actors) { - json ja; - ja["imdb_id"] = actor.imdb_id; - ja["tmdb_id"] = actor.tmdb_id; - ja["jellyfin_id"] = actor.jellyfin_id; - ja["name"] = actor.name; - ja["source_images"] = actor.source_images; - - ja["embeddings"] = json::array(); - for (const auto& emb : actor.embeddings) { - ja["embeddings"].push_back( - std::vector(emb.begin(), emb.end())); - } - j["actors"].push_back(std::move(ja)); + std::string out_path = path; + if (!ends_with(out_path, ".h5") && !ends_with(out_path, ".hdf5")) { + auto dot = out_path.find_last_of('.'); + out_path = (dot == std::string::npos ? out_path : out_path.substr(0, dot)) + ".h5"; + std::cerr << "[gallery] save_gallery: writing HDF5 to " << out_path + << " (galleries are no longer written as JSON)\n"; } - - std::ofstream f(path); - if (!f.is_open()) - throw std::runtime_error("save_gallery: cannot write " + path); - f << j.dump(2) << "\n"; + save_gallery_hdf5(out_path, gallery); } diff --git a/src/gallery/gallery_store.hpp b/src/gallery/gallery_store.hpp index b7c29dc..c038587 100644 --- a/src/gallery/gallery_store.hpp +++ b/src/gallery/gallery_store.hpp @@ -2,9 +2,22 @@ #include "types.hpp" #include -// Load/save the actor gallery from/to a JSON file. +// Load/save the actor gallery. HDF5 (.h5/.hdf5) is the only format written; +// legacy gallery.json files are still readable for backward compatibility but +// save_gallery always writes HDF5 regardless of the requested extension. // -// JSON format: +// HDF5 layout: +// /embeddings float32 [N, 512] all actors' refs concatenated, row-major +// /offset int64 [A] first row of actor a in /embeddings +// /count int32 [A] number of refs for actor a +// /imdb_id /tmdb_id /jellyfin_id /name : variable-length string [A] +// /source_images : variable-length string [N], parallel to /embeddings rows +// /calibration/a, /b : scalar float32 attrs — Platt-sigmoid P(match|sim) fit +// /calibration/valid : scalar int8 attr (0/1) +// /calibration/hash : scalar uint64 attr — hash of the embeddings the fit +// was computed from; a mismatch means "recompute" +// +// Legacy JSON format (read-only): // { // "actors": [ // { diff --git a/src/main.cpp b/src/main.cpp index af9fca0..7e638ae 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -22,7 +22,23 @@ // --extinction actor extinction window in seconds (default: 5.0) // --detector override SCRFD detector model path // --arcface override ArcFace model path +// --scene-detect enable TransNetV2 shot-boundary detection (dense decode; +// writes .scenes.json). Off by default. +// --scene-detector override TransNetV2 .onnx model path +// --scene-detector-engine pre-built TransNetV2 TRT engine (TRT backend) +// --scene-threshold boundary sigmoid prob above this → cut (default: 0.60) +// --scene-stride frames between TransNetV2 windows (default: 50, ≤100) +// --scene-decode-fps dense decode rate in scene-detect mode (default: 12; +// 0 = native fps). Lower = faster, coarser boundaries. +// --dense-scale downscale decoded frames in scene-detect mode (0 max faces kept per frame (default: 10) +// --expand-gallery enable per-film gallery expansion from track continuity +// --expand-buffer per-track diversity buffer size (default: 20) +// --expand-novelty-sim promote only views with best sim < f (default: 0.55) +// --expand-spread-max reject track if buffer spread > f (default: 0.60) +// --expand-min-anchor accepted frames before a track confirms (default: 3) +// --expand-debug-dir

dump promoted mugshots + embeddings here (SAE_DEBUG) // (SAE_DEBUG only) // --debug-dir debug frames output dir (default: debug_frames) // --crop-context bbox expansion factor for context crops (default: 1.5) @@ -31,13 +47,16 @@ #include "types.hpp" #include "gallery/gallery_store.hpp" #include "nodes/frame_source_node.hpp" +#include "nodes/camera_position_change_detector_node.hpp" #include "nodes/face_detector_node.hpp" #include "nodes/face_aligner_node.hpp" #include "nodes/embedder_node.hpp" #include "nodes/face_tracker_node.hpp" #include "nodes/identity_matcher_node.hpp" #include "nodes/scene_tracker_node.hpp" +#include "nodes/scene_detector_node.hpp" #include "nodes/result_sink_node.hpp" +#include "nodes/embedding_dump_node.hpp" #ifdef SAE_DEBUG #include "nodes/debug_renderer_node.hpp" #endif @@ -61,6 +80,7 @@ static Config parse_args(int argc, char** argv) { Config cfg; cfg.detector_model = kDefaultDetectorModel; cfg.arcface_model = kDefaultArcfaceModel; + cfg.scene_model = kDefaultSceneModel; cfg.output_path = "annotations.json"; for (int i = 1; i < argc; ++i) { @@ -73,11 +93,19 @@ static Config parse_args(int argc, char** argv) { if (arg("--movie")) cfg.movie_path = next(); else if (arg("--gallery")) cfg.gallery_path = next(); else if (arg("--output")) cfg.output_path = next(); + else if (arg("--dump-embeddings")) cfg.dump_embeddings_path = next(); else if (arg("--fps")) cfg.sample_fps = std::stof(next()); else if (arg("--max-decode-fps")) cfg.max_decode_fps = std::stof(next()); else if (arg("--start")) cfg.start_sec = std::stod(next()); else if (arg("--end")) cfg.end_sec = std::stod(next()); else if (arg("--cut-threshold")) cfg.cut_threshold = std::stof(next()); + else if (arg("--scene-detect")) cfg.scene_detect = true; + else if (arg("--scene-detector")) cfg.scene_model = next(); + else if (arg("--scene-detector-engine")) cfg.scene_engine = next(); + else if (arg("--scene-threshold")) cfg.scene_threshold = std::stof(next()); + else if (arg("--scene-stride")) cfg.scene_stride = std::stoi(next()); + else if (arg("--scene-decode-fps")) cfg.scene_decode_fps = std::stof(next()); + else if (arg("--dense-scale")) cfg.dense_scale = std::stof(next()); else if (arg("--verbosity")) { int v = std::stoi(next()); cfg.verbosity = v == 2 ? Verbosity::xray : v == 1 ? Verbosity::standard : Verbosity::minimal; } else if (arg("--prior")) cfg.match_prior = std::stof(next()); else if (arg("--prob-threshold")) cfg.prob_threshold = std::stof(next()); @@ -96,7 +124,15 @@ static Config parse_args(int argc, char** argv) { else if (arg("--track-min-iou")) cfg.track_min_iou = std::stof(next()); else if (arg("--track-max-embed")) cfg.track_max_embed_dist = std::stof(next()); else if (arg("--track-max-missing")) cfg.track_max_frames_missing = std::stoi(next()); + else if (arg("--cut-revive-sim")) cfg.cut_revive_sim = std::stof(next()); + else if (arg("--cut-inactive-max")) cfg.cut_inactive_max_frames = std::stoi(next()); else if (arg("--anneal")) cfg.anneal_sec = std::stod(next()); + else if (arg("--expand-gallery")) cfg.expand_gallery = true; + else if (arg("--expand-buffer")) cfg.expand_buffer_size = std::stoi(next()); + else if (arg("--expand-novelty-sim")) cfg.expand_novelty_sim = std::stof(next()); + else if (arg("--expand-spread-max")) cfg.expand_track_spread_max = std::stof(next()); + else if (arg("--expand-min-anchor")) cfg.expand_min_anchor_frames = std::stoi(next()); + else if (arg("--expand-debug-dir")) cfg.expand_debug_dir = next(); else if (arg("--trt-cache")) cfg.trt.cache_dir = next(); else if (arg("--trt-fp16")) cfg.trt.fp16 = true; else if (arg("--no-trt-fp16")) cfg.trt.fp16 = false; @@ -139,16 +175,18 @@ int main(int argc, char** argv) { // ── Construct node functors ─────────────────────────────────────────────── - std::atomic done{false}; + std::atomic done{false}; // set by result_sink (face branch) + std::atomic scene_done{true}; // set by scene_detector; true when disabled - FrameSourceFunc source_fn {cfg}; - FaceDetectorFunc detector_fn{cfg}; - FaceAlignerFunc aligner_fn; - EmbedderFunc embedder_fn{cfg}; - FaceTrackerFunc ftracker_fn{cfg}; - IdentityMatcherFunc matcher_fn {gallery, cfg}; - SceneTrackerFunc tracker_fn {cfg}; - ResultSinkFunc sink_fn {cfg, done}; + FrameSourceFunc source_fn {cfg}; + CameraPositionChangeDetectorFunc campos_fn {cfg}; + FaceDetectorFunc detector_fn{cfg}; + FaceAlignerFunc aligner_fn; + EmbedderFunc embedder_fn{cfg}; + FaceTrackerFunc ftracker_fn{cfg}; + IdentityMatcherFunc matcher_fn {gallery, cfg}; + SceneTrackerFunc tracker_fn {cfg}; + ResultSinkFunc sink_fn {cfg, done}; #ifdef SAE_DEBUG DebugRendererFunc debug_fn {cfg}; #endif @@ -158,7 +196,8 @@ int main(int argc, char** argv) { // Queue sizes tuned to the pipeline's speed profile: // embedder (16ms) is the slowest GPU node — buffer before it must be largest // to prevent face_aligner pool overflows and frame drops. - kpn::ObjectNode, kpn::out<"frame">, "frame_source", 0> source (source_fn, 32); + kpn::ObjectNode, kpn::out<"raw">, "frame_source", 0> source (source_fn, 32); + kpn::ObjectNode, kpn::out<"frame">, "camera_pos", 0> campos (campos_fn, 32); kpn::ObjectNode, kpn::out<"scene">, "face_detector", 0> detector (detector_fn, 64); kpn::ObjectNode, kpn::out<"aligned">, "face_aligner", 0> aligner (aligner_fn, 64); kpn::ObjectNode, kpn::out<"embedded">, "embedder", 0> embedder (embedder_fn, 32); @@ -167,81 +206,143 @@ int main(int argc, char** argv) { kpn::ObjectNode, kpn::out<"annotation">, "scene_tracker", 0> tracker (tracker_fn, 16); kpn::ObjectNode,kpn::out<>, "result_sink", 0> sink (sink_fn, 16); - // ── Build static network ────────────────────────────────────────────────── - -#ifdef SAE_DEBUG - kpn::ObjectNode, kpn::out<>, "debug_renderer", 1> debug_node(debug_fn, 16); - - // matcher → FanoutNode → scene_tracker + debug_node (auto-inserted) - auto net = kpn::make_network( - kpn::edge(source.output<"frame">(), detector.input<"frame">()), - kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), - kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), - kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()), - kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()), - kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()), - kpn::edge(matcher.output<"matched">(), debug_node.input<"matched">()), - kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">()) - ); -#else - auto net = kpn::make_network( - kpn::edge(source.output<"frame">(), detector.input<"frame">()), - kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), - kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), - kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()), - kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()), - kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()), - kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">()) - ); -#endif - - // ── Pipeline observability (KPN event handler) ──────────────────────────── - // Tally dropped frames per node (channel overflow) and detect a node that - // stops unexpectedly. A Closed event from any node other than result_sink at - // EOF means a stage crashed — without this the main loop below would hang on - // `done` forever, so we trip a flag to unblock it and exit non-zero. + // ── Pipeline observability + run loop (topology-agnostic) ────────────────── + // Factored into a lambda so the two topologies (with/without the scene-detect + // branch) share identical event handling, wait loop, and teardown. Any + // make_network result type binds to `Net&&`. std::mutex event_mtx; std::map overflow_counts; std::atomic node_crashed{false}; - net.set_event_handler( - [&](std::string_view node_name, kpn::NodeEvent ev, - std::chrono::steady_clock::time_point) { - if (ev == kpn::NodeEvent::Overflow) { - std::lock_guard lk(event_mtx); - ++overflow_counts[std::string(node_name)]; - } else { // NodeEvent::Closed - // result_sink closing once EOF has been signalled is the normal - // shutdown path, not a crash. - if (node_name == "result_sink" && done.load(std::memory_order_acquire)) - return; - std::cerr << "[main] node '" << node_name - << "' stopped unexpectedly — aborting pipeline\n"; - node_crashed.store(true, std::memory_order_release); + auto run_net = [&](auto&& net) -> int { + // Tally per-node channel overflow, and treat any non-result_sink Closed + // event as a crash so the wait loop below can't hang on `done` forever. + net.set_event_handler( + [&](std::string_view node_name, kpn::NodeEvent ev, + std::chrono::steady_clock::time_point) { + if (ev == kpn::NodeEvent::Overflow) { + std::lock_guard lk(event_mtx); + ++overflow_counts[std::string(node_name)]; + } else { // NodeEvent::Closed + if (node_name == "result_sink" && done.load(std::memory_order_acquire)) + return; + std::cerr << "[main] node '" << node_name + << "' stopped unexpectedly — aborting pipeline\n"; + node_crashed.store(true, std::memory_order_release); + } + }); + + std::cerr << "[main] starting pipeline…\n"; + net.start(); + + // Wait until BOTH terminal branches finish: result_sink (face pipeline) + // and, when enabled, scene_detector (the dense TransNetV2 branch, which + // runs much slower and must not be torn down mid-stream). scene_done is + // pre-set true when scene detection is disabled. + while ((!done.load(std::memory_order_acquire) || + !scene_done.load(std::memory_order_acquire)) && + !node_crashed.load(std::memory_order_acquire)) + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + net.stop(); + net.print_diagnostics(); + + { + std::lock_guard lk(event_mtx); + if (!overflow_counts.empty()) { + std::cerr << "[main] dropped frames (channel overflow):\n"; + for (const auto& [name, count] : overflow_counts) + std::cerr << " " << name << ": " << count << "\n"; } - }); - - // ── Run ─────────────────────────────────────────────────────────────────── - std::cerr << "[main] starting pipeline…\n"; - net.start(); - - // Main thread waits until ResultSinkFunc signals EOF completion, or a node - // crash trips node_crashed. - while (!done.load(std::memory_order_acquire) && - !node_crashed.load(std::memory_order_acquire)) - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - net.stop(); - net.print_diagnostics(); - - { - std::lock_guard lk(event_mtx); - if (!overflow_counts.empty()) { - std::cerr << "[main] dropped frames (channel overflow):\n"; - for (const auto& [name, count] : overflow_counts) - std::cerr << " " << name << ": " << count << "\n"; } - } + return node_crashed.load(std::memory_order_acquire) ? 1 : 0; + }; - return node_crashed.load(std::memory_order_acquire) ? 1 : 0; + // ── Build static network and run ────────────────────────────────────────── + // Common face-analysis chain (campos → … → sink) is identical in all cases; + // the scene-detect branch and the debug fanout are spliced on conditionally. + // Topology: + // plain: source → campos → detector → … → sink + // scene-detect: source ─┬→ campos → decimate(filter) → detector → … → sink + // └→ scene_detector (TransNetV2 sink → scenes.json) + // The fanout after `source` is auto-inserted by make_network when its output + // feeds two edges. In dense mode campos still sees native-rate frames (so it + // detects angle changes correctly); a FilterNode then thins to sample_fps + // before face detection. +#ifdef SAE_DEBUG + kpn::ObjectNode, kpn::out<>, "debug_renderer", 1> debug_node(debug_fn, 16); + #define SAE_DEBUG_EDGE , kpn::edge(matcher.output<"matched">(), debug_node.input<"matched">()) +#else + #define SAE_DEBUG_EDGE +#endif + + int rc = 0; + if (!cfg.dump_embeddings_path.empty()) { + // Dump-only topology: run the expensive front half and tee the embedder + // output to an HDF5 dump for offline sweep replay (sae_kpn). Downstream + // matching is skipped — the sweep re-runs it from the dump. + EmbeddingDumpFunc dump_fn{cfg, done}; + kpn::ObjectNode, kpn::out<>, "embedding_dump", 0> + dump_node(dump_fn, 32); + auto net = kpn::make_network( + kpn::edge(source.output<"raw">(), campos.input<"raw">()), + kpn::edge(campos.output<"frame">(), detector.input<"frame">()), + kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), + kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), + kpn::edge(embedder.output<"embedded">(), dump_node.input<"embedded">()) + ); + return run_net(std::move(net)); + } + if (cfg.scene_detect) { + scene_done.store(false, std::memory_order_release); // now a real terminal branch + SceneDetectorFunc scene_fn{cfg, scene_done}; + kpn::ObjectNode, kpn::out<>, "scene_detector", 0> + scene_node(scene_fn, 128); + + // Decimator: keep frames on the sample_fps cadence, drop the rest. + // eof always passes so downstream shuts down cleanly. Stateful — one + // instance, mutable via shared_ptr so the std::function stays copyable. + auto decim_state = std::make_shared(-1e18); + const double interval = 1.0 / cfg.sample_fps; + auto decimate = kpn::make_filter( + [decim_state, interval](const Frame& f) { + if (f.eof) return true; + if (f.timestamp_sec - *decim_state >= interval - 1e-6) { + *decim_state = f.timestamp_sec; + return true; + } + return false; + }, 32); + + auto net = kpn::make_network( + kpn::edge(source.output<"raw">(), campos.input<"raw">()), + kpn::edge(source.output<"raw">(), scene_node.input<"dense">()), + kpn::edge(campos.output<"frame">(), decimate.input<0>()), + kpn::edge(decimate.output<0>(), detector.input<"frame">()), + kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), + kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), + kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()), + kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()), + kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()), + kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">()) + SAE_DEBUG_EDGE + ); + rc = run_net(std::move(net)); + } else { + auto net = kpn::make_network( + kpn::edge(source.output<"raw">(), campos.input<"raw">()), + kpn::edge(campos.output<"frame">(), detector.input<"frame">()), + kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), + kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), + kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()), + kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()), + kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()), + kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">()) + SAE_DEBUG_EDGE + ); + rc = run_net(std::move(net)); + } +#undef SAE_DEBUG_EDGE + + return rc; } diff --git a/src/nodes/face_detector_node.hpp b/src/nodes/face_detector_node.hpp index 2f5b69a..0cb2ddf 100644 --- a/src/nodes/face_detector_node.hpp +++ b/src/nodes/face_detector_node.hpp @@ -26,10 +26,15 @@ struct FaceDetectorFunc { auto faces = detector_->detect(f.image); - // Drop faces below minimum pixel size (too small for reliable ArcFace alignment) + // Drop faces below minimum pixel size (too small for reliable ArcFace + // alignment). Note: when dense_scale downscaled the frame, both the + // detection coords and min_face_px are in downscaled space — so scale + // the threshold down to match, keeping the physical size cutoff constant. + const float min_px = (f.bbox_upscale != 1.f) + ? min_face_px_ / f.bbox_upscale : min_face_px_; faces.erase( std::remove_if(faces.begin(), faces.end(), [&](const DetectedFace& d) { - return d.bbox.width < min_face_px_ || d.bbox.height < min_face_px_; + return d.bbox.width < min_px || d.bbox.height < min_px; }), faces.end()); diff --git a/src/nodes/face_tracker_node.hpp b/src/nodes/face_tracker_node.hpp index 61f5c02..5a73a02 100644 --- a/src/nodes/face_tracker_node.hpp +++ b/src/nodes/face_tracker_node.hpp @@ -23,6 +23,16 @@ // // Unmatched tracks have their frames_missing counter incremented; they are // expired once frames_missing > max_frames_missing. +// +// Cross-cut re-association. A camera-angle change (Frame::is_cut, set by +// camera_position_change_detector) destroys spatial (IoU) continuity — the same +// person reappears at a new position — but not identity. On a cut the tracker +// does NOT discard its tracks; it parks them in an inactive pool keyed by their +// last-frame raw embedding. A post-cut detection whose raw cosine similarity to +// a parked track's last-frame embedding is ≥ cut_revive_sim revives that track: +// the original track_id, mean embedding and n_frames are restored (only the bbox +// jumps to the new detection), so identity continuity survives the cut. Parked +// tracks left unrevived for cut_inactive_max_frames are finally dropped. struct FaceTrackerFunc { static constexpr std::string_view label() { return "face_tracker"; } @@ -30,6 +40,7 @@ struct FaceTrackerFunc { struct TrackState { cv::Rect2f bbox; Embedding mean_emb{}; + Embedding last_emb{}; // raw embedding of the most recent matched frame int n_frames{0}; int frames_missing{0}; }; @@ -39,16 +50,21 @@ struct FaceTrackerFunc { , min_iou_(cfg.track_min_iou) , max_embed_dist_(cfg.track_max_embed_dist) , max_missing_(cfg.track_max_frames_missing) + , revive_sim_(cfg.cut_revive_sim) + , inactive_max_(cfg.cut_inactive_max_frames) { std::cerr << "[face_tracker] alpha=" << alpha_ << " min_iou=" << min_iou_ << " max_embed_dist=" << max_embed_dist_ - << " max_missing=" << max_missing_ << "\n"; + << " max_missing=" << max_missing_ + << " cut_revive_sim=" << revive_sim_ + << " cut_inactive_max=" << inactive_max_ << "\n"; } TrackedSceneFrame operator()(EmbeddedSceneFrame ef) { if (ef.source.eof) { tracks_.clear(); + inactive_.clear(); TrackedSceneFrame out; out.source = std::move(ef.source); return out; @@ -56,11 +72,26 @@ struct FaceTrackerFunc { const int n_det = static_cast(ef.embeddings.size()); + // Camera-angle change: park active tracks instead of destroying them so + // they can be revived by identity (raw last-frame embedding cosine) once + // the same people reappear from the new angle. if (ef.source.is_cut && !tracks_.empty()) { - std::cerr << "[face_tracker] cut — clearing " << tracks_.size() << " tracks\n"; + std::cerr << "[face_tracker] cut — parking " << tracks_.size() + << " track(s) into inactive pool\n"; + for (auto& [tid, ts] : tracks_) { + ts.frames_missing = 0; // repurpose as time-since-parked counter + inactive_[tid] = std::move(ts); + } tracks_.clear(); } + // Age the inactive pool every frame and drop tracks parked too long. + for (auto it = inactive_.begin(); it != inactive_.end(); ) { + it->second.frames_missing++; + it = (it->second.frames_missing > inactive_max_) + ? inactive_.erase(it) : std::next(it); + } + // Snapshot active track IDs so the map can be modified safely below std::vector tids; tids.reserve(tracks_.size()); @@ -111,6 +142,7 @@ struct FaceTrackerFunc { continue; } update_mean(ts.mean_emb, ts.n_frames, ef.embeddings[di]); + ts.last_emb = ef.embeddings[di]; ts.bbox = ef.faces[di].bbox; ts.n_frames++; ts.frames_missing = 0; @@ -119,13 +151,34 @@ struct FaceTrackerFunc { out.track_ids[di] = tids[ti]; } - // Create new tracks for unmatched detections + // Handle unmatched detections: first try to revive a parked track by + // identity (raw last-frame embedding cosine), else start a fresh track. for (int di = 0; di < n_det; ++di) { if (det_matched[di]) continue; - int tid = next_id_++; + + int tid = revive_from_inactive(ef.embeddings[di]); + if (tid >= 0) { + // Restore the parked track: keep its identity statistics + // (mean_emb, n_frames), jump the bbox to the new detection. + TrackState ts = std::move(inactive_[tid]); + inactive_.erase(tid); + update_mean(ts.mean_emb, ts.n_frames, ef.embeddings[di]); + ts.last_emb = ef.embeddings[di]; + ts.bbox = ef.faces[di].bbox; + ts.n_frames++; + ts.frames_missing = 0; + tracks_[tid] = std::move(ts); + out.track_ids[di] = tid; + std::cerr << "[face_tracker] revived track " << tid + << " across cut\n"; + continue; + } + + tid = next_id_++; TrackState ts; ts.bbox = ef.faces[di].bbox; ts.mean_emb = ef.embeddings[di]; + ts.last_emb = ef.embeddings[di]; ts.n_frames = 1; tracks_[tid] = ts; out.track_ids[di] = tid; @@ -141,6 +194,21 @@ struct FaceTrackerFunc { } private: + // Pick the parked track whose last-frame embedding is most similar to emb, + // returning its id if that raw cosine similarity clears revive_sim_, else -1. + // The caller removes the returned track from the pool, so a later detection in + // the same frame cannot claim it again. + int revive_from_inactive(const Embedding& emb) const { + int best_tid = -1; + float best_sim = revive_sim_; // threshold is the bar to beat (inclusive) + for (const auto& [tid, ts] : inactive_) { + float sim = cosine_similarity(ts.last_emb, emb); + if (sim >= best_sim) { best_sim = sim; best_tid = tid; } + // subsequent ties keep the later id; harmless, all clear the threshold + } + return best_tid; + } + // IoU of two axis-aligned bounding boxes static float iou(const cv::Rect2f& a, const cv::Rect2f& b) { float ix = std::max(0.f, std::min(a.x + a.width, b.x + b.width) @@ -225,9 +293,12 @@ private: } std::map tracks_; + std::map inactive_; // parked across a cut, keyed by track id int next_id_{0}; float alpha_; float min_iou_; float max_embed_dist_; int max_missing_; + float revive_sim_; + int inactive_max_; }; diff --git a/src/nodes/frame_source_node.hpp b/src/nodes/frame_source_node.hpp index ae66c49..58c452d 100644 --- a/src/nodes/frame_source_node.hpp +++ b/src/nodes/frame_source_node.hpp @@ -3,7 +3,6 @@ #include "config.hpp" #include "ffmpeg_decoder.hpp" -#include #include #include #include @@ -27,21 +26,46 @@ struct FrameSourceFunc { static constexpr std::string_view label() { return "frame_source"; } explicit FrameSourceFunc(const Config& cfg) - : decoder_(std::make_unique(cfg.movie_path)) + : decoder_(std::make_unique( + cfg.movie_path, /*use_hw=*/true, + /*out_scale=*/cfg.scene_detect ? cfg.dense_scale : 1.0f)) { sample_interval_sec_ = 1.0 / cfg.sample_fps; next_pos_sec_ = cfg.start_sec; end_sec_ = cfg.end_sec; - cut_threshold_ = cfg.cut_threshold; max_decode_fps_ = cfg.max_decode_fps; + // Dense mode: emit every native-rate frame instead of seeking to each + // 1-FPS sample point. Required by the TransNetV2 scene detector, which + // needs consecutive frames. A downstream decimator drops back to + // sample_fps for the face pipeline. When dense, we advance by the + // decoder's frame period (best effort — read_at decodes forward past the + // last position, so consecutive small steps yield consecutive frames). + dense_ = cfg.scene_detect; + double dense_fps = 0.0; + if (dense_) { + double vfps = decoder_->fps(); + if (vfps <= 0.0) vfps = 25.0; + // Dense decode rate: capped at native fps. A lower scene_decode_fps + // decodes fewer frames (big speedup); TransNetV2 still localises cuts + // and boundary timestamps stay exact (keyed off real timestamps). + dense_fps = (cfg.scene_decode_fps > 0.f) + ? std::min(static_cast(cfg.scene_decode_fps), vfps) + : vfps; + dense_step_sec_ = 1.0 / dense_fps; + if (cfg.dense_scale > 0.f && cfg.dense_scale < 1.f) + bbox_upscale_ = 1.0f / cfg.dense_scale; + } + double total_s = decoder_->duration_sec(); double span_s = (end_sec_ > 0 ? std::min(end_sec_, total_s) : total_s) - cfg.start_sec; - int n_frames = static_cast(span_s * cfg.sample_fps); + double emit_fps = dense_ ? dense_fps : cfg.sample_fps; + int n_frames = static_cast(span_s * emit_fps); std::cerr << "[frame_source] decoder=" << decoder_->codec_name() << " (" << decoder_->hw_backend() << ")" << " video_fps=" << decoder_->fps() + << (dense_ ? " DENSE@" + std::to_string(dense_fps) + "fps" : "") << " start=" << cfg.start_sec << "s" << (end_sec_ > 0 ? " end=" + std::to_string(end_sec_) + "s" : "") << " sample_fps=" << cfg.sample_fps @@ -97,29 +121,14 @@ struct FrameSourceFunc { return Frame{{}, next_pos_sec_, frame_idx_++, /*eof=*/true}; } - // Cut detection: compare grayscale histogram to previous frame - bool is_cut = false; - cv::Mat gray; - cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY); - cv::Mat hist; - const int bins = 64; - const float range[] = {0.f, 256.f}; - const float* ranges = range; - cv::calcHist(&gray, 1, nullptr, cv::Mat(), hist, 1, &bins, &ranges); - cv::normalize(hist, hist, 1.0, 0.0, cv::NORM_L1); - - if (prev_hist_valid_) { - double corr = cv::compareHist(prev_hist_, hist, cv::HISTCMP_CORREL); - is_cut = (corr < cut_threshold_); - if (is_cut) - std::cerr << "[frame_source] cut at t=" << next_pos_sec_ - << "s hist_corr=" << corr << "\n"; - } - prev_hist_ = hist; - prev_hist_valid_ = true; - - Frame f{img, next_pos_sec_, frame_idx_++, /*eof=*/false, is_cut}; - next_pos_sec_ += sample_interval_sec_; + // Cut detection is a downstream concern: camera_position_change_detector + // owns the histogram compare and sets Frame::is_cut. The source emits + // is_cut=false and only decodes/samples frames. + Frame f{img, next_pos_sec_, frame_idx_++, /*eof=*/false, /*is_cut=*/false}; + // When dense_scale downscaled the decode, tell downstream how to map + // detector coordinates back to original video resolution. + f.bbox_upscale = bbox_upscale_; + next_pos_sec_ += dense_ ? dense_step_sec_ : sample_interval_sec_; if (end_sec_ > 0 && next_pos_sec_ > end_sec_) { hit_eof_ = true; std::cerr << "[frame_source] reached end_sec=" << end_sec_ << "s\n"; @@ -130,16 +139,16 @@ struct FrameSourceFunc { private: std::unique_ptr decoder_; double sample_interval_sec_{1.0}; + bool dense_{false}; + float bbox_upscale_{1.f}; + double dense_step_sec_{0.04}; double next_pos_sec_{0.0}; double end_sec_{-1.0}; - float cut_threshold_{0.70f}; float max_decode_fps_{0.f}; std::chrono::steady_clock::time_point next_decode_at_{}; bool rate_started_{false}; int64_t frame_idx_{0}; bool hit_eof_{false}; - cv::Mat prev_hist_; - bool prev_hist_valid_{false}; double decode_ms_acc_{0.0}; int decode_count_{0}; }; diff --git a/src/nodes/identity_matcher_node.hpp b/src/nodes/identity_matcher_node.hpp index 6da9588..7738995 100644 --- a/src/nodes/identity_matcher_node.hpp +++ b/src/nodes/identity_matcher_node.hpp @@ -4,6 +4,7 @@ #include "inference/similarity.hpp" #include "gallery/gallery_store.hpp" #include "gallery/gallery_calibration.hpp" +#include "gallery/track_gallery.hpp" #include #include @@ -51,6 +52,7 @@ struct IdentityMatcherFunc { , threshold_(cfg.match_threshold) , ratio_(cfg.match_ratio) , ratio_ceil_(cfg.match_ratio_ceil) + , track_gallery_(cfg) { std::cerr << "[identity_matcher] flattening gallery embeddings...\n"; for (int ai = 0; ai < static_cast(gallery_.actors.size()); ++ai) { @@ -63,8 +65,24 @@ struct IdentityMatcherFunc { std::cerr << "[identity_matcher] starting calibration (" << flat_emb_.size() << " embeddings)...\n"; + bool recomputed = false; cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_, - cfg.gallery_path + ".calib_cache.json"); + gallery_.calib_a, gallery_.calib_b, + gallery_.calib_valid, gallery_.calib_hash, + cfg.gallery_path + ".calib_cache", recomputed); + if (recomputed) { + // Persist the freshly-fitted calibration into the gallery file (always + // HDF5 — save_gallery rewrites any other extension, see gallery_store.cpp) + // so the next run against this same, unchanged gallery skips the O(n^2) fit. + ActorGallery to_save = gallery_; + to_save.calib_a = cal_.a; + to_save.calib_b = cal_.b; + to_save.calib_valid = cal_.valid; + to_save.calib_hash = hash_gallery_embeddings(flat_emb_, flat_actor_); + save_gallery(cfg.gallery_path, to_save); + std::cerr << "[identity_matcher] wrote refreshed calibration back to " + << cfg.gallery_path << "\n"; + } if (cal_.valid) { std::cerr << "[identity_matcher] calibrated Bayesian matching" @@ -89,8 +107,23 @@ struct IdentityMatcherFunc { sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces); } + // Runtime setter — lets a persistent pipeline be reused across a threshold sweep + // without rebuilding the (expensive, gallery-resident) matcher. The gallery, + // calibration and GPU sim-engine stay put; only the accept threshold changes. + void set_prob_threshold(float t) { prob_threshold_ = t; } + MatchedSceneFrame operator()(TrackedSceneFrame tf) { - if (tf.source.eof) return {std::move(tf.source), {}}; + if (tf.source.eof) { + track_gallery_.clear_tracks(); + return {std::move(tf.source), {}}; + } + + // A hard cut changes the camera viewpoint. The face_tracker may revive a + // track_id across the cut (identity continuity), but promotion must never + // mix embeddings from two viewpoints under one buffer, so we still drop + // every diversity buffer here — a revived track simply re-accumulates its + // buffer from post-cut frames. Stale cross-cut embeddings are never promoted. + if (tf.source.is_cut) track_gallery_.clear_tracks(); const int n_faces = static_cast(tf.embeddings.size()); std::vector actors; @@ -120,6 +153,14 @@ struct IdentityMatcherFunc { if (sim > best_sim[ai]) best_sim[ai] = sim; } + // Fold in the per-film annex (CPU-side, tens of embeddings). Promoted + // pose-varied views compete for best-of-N exactly like baked refs, so + // a face at a pose the gallery lacked can now win its true actor. + for (const auto& ae : track_gallery_.annex()) { + float sim = cosine_similarity(tf.embeddings[fi], ae.emb); + if (sim > best_sim[ae.actor_idx]) best_sim[ae.actor_idx] = sim; + } + int best_actor = -1; int second_actor = -1; float best_s = -std::numeric_limits::max(); @@ -155,7 +196,15 @@ struct IdentityMatcherFunc { } IdentifiedActor ia; + // Map bbox back to original video resolution when dense_scale + // downscaled the decoded frame (detection/tracking ran downscaled; + // output bboxes must be in original pixel space). ia.bbox = tf.faces[fi].bbox; + if (tf.source.bbox_upscale != 1.f) { + const float s = tf.source.bbox_upscale; + ia.bbox.x *= s; ia.bbox.y *= s; + ia.bbox.width *= s; ia.bbox.height *= s; + } ia.crop = tf.crops[fi]; ia.track_id = tf.track_ids[fi]; @@ -170,6 +219,14 @@ struct IdentityMatcherFunc { : best_s; } + // Feed this face into per-film gallery expansion. best_actor/best_s + // reflect the actor with the strongest gallery similarity for this + // face (annex already folded in above); the track's diversity buffer + // keeps the gallery-far views and promotes them once the track is + // confirmed. No-op unless --expand-gallery is set. + track_gallery_.observe(tf.track_ids[fi], tf.embeddings[fi], + best_actor, best_s, accept, tf.crops[fi]); + actors.push_back(std::move(ia)); } @@ -189,4 +246,5 @@ private: int n_gallery_{0}; std::unique_ptr sim_engine_; + TrackGallery track_gallery_; }; diff --git a/src/nodes/preview_node.hpp b/src/nodes/preview_node.hpp index 53d6bf9..ebe0133 100644 --- a/src/nodes/preview_node.hpp +++ b/src/nodes/preview_node.hpp @@ -7,6 +7,7 @@ #include #include +#include #include // ── PreviewNode ─────────────────────────────────────────────────────────────── @@ -35,6 +36,7 @@ public: cv::Mat display = mf.source.image.clone(); draw_detections(display, mf.actors); draw_hud(display, mf.source.timestamp_sec, mf.actors); + draw_cut_meter(display, mf.source.cut_score, mf.source.is_cut); // Fit to display width while keeping aspect ratio if (display.cols > max_display_w_) { @@ -132,6 +134,45 @@ private: } } + // Running cut-score meter (top-right): a horizontal bar whose fill and colour + // track the histogram cut score in [0,1] (0 = no change, ~1 = hard cut). + // Flashes a red "CUT" tag on frames where the threshold tripped (is_cut). + static void draw_cut_meter(cv::Mat& img, float score, bool is_cut) { + const int bar_w = 220; + const int bar_h = 14; + const int pad = 10; + const int x0 = img.cols - bar_w - pad; + const int y0 = pad + 16; // below any top-left HUD row height + + score = std::clamp(score, 0.f, 1.f); + + // Label + char lbl[32]; + std::snprintf(lbl, sizeof(lbl), "cut %.2f", score); + cv::putText(img, lbl, cv::Point(x0, y0 - 4), + cv::FONT_HERSHEY_SIMPLEX, 0.5, + cv::Scalar(220, 220, 220), 1, cv::LINE_AA); + + // Track (dark) + fill (green→red by score) + cv::rectangle(img, cv::Rect(x0, y0, bar_w, bar_h), + cv::Scalar(40, 40, 40), cv::FILLED); + int fill_w = static_cast(bar_w * score); + // BGR: low score → green (0,210,60), high → red (0,0,255) + cv::Scalar fill_col(60.0 * (1.f - score), + 210.0 * (1.f - score), + 60.0 * (1.f - score) + 255.0 * score); + if (fill_w > 0) + cv::rectangle(img, cv::Rect(x0, y0, fill_w, bar_h), fill_col, cv::FILLED); + cv::rectangle(img, cv::Rect(x0, y0, bar_w, bar_h), + cv::Scalar(120, 120, 120), 1); + + if (is_cut) { + cv::putText(img, "CUT", cv::Point(x0 - 52, y0 + bar_h), + cv::FONT_HERSHEY_DUPLEX, 0.6, + cv::Scalar(0, 0, 255), 2, cv::LINE_AA); + } + } + static std::string pct(float v) { char buf[8]; std::snprintf(buf, sizeof(buf), "%.0f%%", v * 100.f); diff --git a/src/nodes/scene_tracker_node.hpp b/src/nodes/scene_tracker_node.hpp index 5ca6e31..2de6897 100644 --- a/src/nodes/scene_tracker_node.hpp +++ b/src/nodes/scene_tracker_node.hpp @@ -26,6 +26,10 @@ struct SceneTrackerFunc { std::cerr << "[scene_tracker] extinction_sec=" << extinction_sec_ << "\n"; } + // Runtime setter for pipeline reuse across a sweep. Also clears the active-actor + // state so a re-run starts clean (no carry-over from the previous config's film). + void set_extinction_sec(double s) { extinction_sec_ = s; active_.clear(); } + SceneAnnotation operator()(MatchedSceneFrame mf) { if (mf.source.eof) return {0.0, {}, /*eof=*/true}; diff --git a/src/scene_preview.cpp b/src/scene_preview.cpp index 7019854..b1a3e27 100644 --- a/src/scene_preview.cpp +++ b/src/scene_preview.cpp @@ -3,8 +3,11 @@ // // KPN topology: // -// [frame_source] ──► [face_detector] ──► [face_aligner] ──► [embedder] +// [frame_source] ──► [camera_pos] ──► [face_detector] ──► [face_aligner] ──► [embedder] // ──► [identity_matcher] ──► FanoutNode +// +// camera_pos (histogram cut detector) stamps Frame::cut_score / is_cut, which +// ride through to the preview HUD's cut-score meter. // ├──► [scene_tracker] ──► [result_sink] (background thread) // └──► [preview_node] (main thread) // @@ -20,6 +23,7 @@ #include "types.hpp" #include "gallery/gallery_store.hpp" #include "nodes/frame_source_node.hpp" +#include "nodes/camera_position_change_detector_node.hpp" #include "nodes/face_detector_node.hpp" #include "nodes/face_aligner_node.hpp" #include "nodes/embedder_node.hpp" @@ -87,6 +91,25 @@ static Config parse_args(int argc, char** argv) { else if (arg("--no-trt-fp16")) cfg.trt.fp16 = false; else if (arg("--trt-int8")) cfg.trt.int8 = true; else if (arg("--embed-batch")) cfg.embed_batch_size = std::stoi(next()); + // Per-film gallery expansion — preview supports it (same cfg fields). + else if (arg("--expand-gallery")) cfg.expand_gallery = true; + else if (arg("--expand-buffer")) cfg.expand_buffer_size = std::stoi(next()); + else if (arg("--expand-novelty-sim")) cfg.expand_novelty_sim = std::stof(next()); + else if (arg("--expand-spread-max")) cfg.expand_track_spread_max = std::stof(next()); + else if (arg("--expand-min-anchor")) cfg.expand_min_anchor_frames = std::stoi(next()); + // Scene detection is scene_analyze-only (needs the dense TransNetV2 branch). + // Accept the flags so a shared command line runs, but note they're inert + // here — the preview shows the histogram cut-score meter instead. + else if (arg("--scene-detect")) { + std::cerr << "[preview] note: --scene-detect is inert in preview " + "(TransNetV2 needs the dense scene_analyze pipeline); " + "showing the histogram cut-score meter instead\n"; + } + else if (arg("--scene-detector") || arg("--scene-detector-engine") || + arg("--scene-threshold") || arg("--scene-stride") || + arg("--scene-decode-fps") || arg("--dense-scale")) { + next(); // consume the value; inert in preview + } else { std::cerr << "[warn] unknown flag: " << argv[i] << "\n"; } } if (cfg.movie_path.empty()) throw std::runtime_error("--movie is required"); @@ -112,7 +135,8 @@ int main(int argc, char** argv) { // ── Functors ────────────────────────────────────────────────────────────── std::atomic done{false}; - FrameSourceFunc source_fn {cfg}; + FrameSourceFunc source_fn {cfg}; + CameraPositionChangeDetectorFunc campos_fn {cfg}; FaceDetectorFunc detector_fn{cfg}; FaceAlignerFunc aligner_fn; EmbedderFunc embedder_fn{cfg}; @@ -122,7 +146,8 @@ int main(int argc, char** argv) { ResultSinkFunc sink_fn {cfg, done}; // ── KPN ObjectNodes ─────────────────────────────────────────────────────── - kpn::ObjectNode, kpn::out<"frame">, "frame_source", 0> source (source_fn, 32); + kpn::ObjectNode, kpn::out<"raw">, "frame_source", 0> source (source_fn, 32); + kpn::ObjectNode, kpn::out<"frame">, "camera_pos", 0> campos (campos_fn, 32); kpn::ObjectNode, kpn::out<"scene">, "face_detector", 0> detector (detector_fn, 64); kpn::ObjectNode, kpn::out<"aligned">, "face_aligner", 0> aligner (aligner_fn, 64); kpn::ObjectNode, kpn::out<"embedded">, "embedder", 0> embedder (embedder_fn, 32); @@ -136,7 +161,8 @@ int main(int argc, char** argv) { // matcher → FanoutNode → [scene_tracker, preview] (auto-inserted) auto net = kpn::make_network( - kpn::edge(source.output<"frame">(), detector.input<"frame">()), + kpn::edge(source.output<"raw">(), campos.input<"raw">()), + kpn::edge(campos.output<"frame">(), detector.input<"frame">()), kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()), diff --git a/src/types.hpp b/src/types.hpp index 38c13ed..08afb42 100644 --- a/src/types.hpp +++ b/src/types.hpp @@ -24,7 +24,20 @@ struct Frame { double timestamp_sec{0.0}; int64_t frame_idx{-1}; bool eof{false}; - bool is_cut{false}; // true when a hard scene cut was detected before this frame + bool is_cut{false}; // histogram: intra-scene camera-angle change (tracker reset) + bool is_scene_boundary{false}; // TransNetV2: true shot/scene boundary (opt-in) + float cut_score{0.f}; // histogram cut score = 1 - hist_corr (0=identical, ~1=cut); HUD/debug + float bbox_upscale{1.f}; // multiply detector bboxes/landmarks by this to map back to + // original video resolution (>1 when dense_scale downscaled the frame) +}; + +// ── CutEvent ────────────────────────────────────────────────────────────────── +// Emitted by SceneDetectorFunc when TransNetV2 localises a shot boundary, keyed +// by the boundary frame's timestamp. eof=true is the shutdown sentinel. +struct CutEvent { + double timestamp_sec{0.0}; + float probability{0.f}; // sigmoid boundary score at the peak + bool eof{false}; }; // ── ArcFace alignment ───────────────────────────────────────────────────────── @@ -122,4 +135,13 @@ struct ActorGallery { std::vector source_images; }; std::vector actors; + + // Cached Platt-sigmoid calibration (see gallery/gallery_calibration.hpp), + // stored alongside the gallery in HDF5 so it never needs recomputing + // unless the reference embeddings actually change. calib_valid=false and + // calib_hash=0 means "not present in this file, compute it." + float calib_a{10.f}; + float calib_b{-5.f}; + bool calib_valid{false}; + uint64_t calib_hash{0}; };