From a3ba53ddf7b421bf2193e452de968cf2ea2fb552 Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sat, 13 Jun 2026 22:44:44 +0200 Subject: [PATCH] improved performance --- CMakeLists.txt | 11 ++- README.md | 9 +- src/config.hpp | 1 - src/main.cpp | 1 - src/nodes/face_tracker_node.hpp | 18 +--- src/nodes/identity_matcher_node.hpp | 138 +++++++++++++++++++++++++--- src/scene_preview.cpp | 1 - src/types.hpp | 5 +- 8 files changed, 142 insertions(+), 42 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ecf078..e5b0ea5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,18 +44,21 @@ find_library(CUDART_LIB cudart find_path(CUDART_INCLUDE cuda_runtime_api.h HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include /usr/local/cuda/include /usr/include) -if(NOT (NVINFER_LIB AND NVINFER_INCLUDE AND CUDART_LIB AND CUDART_INCLUDE)) +find_library(CUBLAS_LIB cublas + HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64 + /usr/local/cuda/lib64 /usr/lib) +if(NOT (NVINFER_LIB AND NVINFER_INCLUDE AND CUDART_LIB AND CUDART_INCLUDE AND CUBLAS_LIB)) message(FATAL_ERROR "TensorRT or CUDA runtime not found " "(nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE} " - "cudart=${CUDART_LIB} headers=${CUDART_INCLUDE})") + "cudart=${CUDART_LIB} cublas=${CUBLAS_LIB} headers=${CUDART_INCLUDE})") endif() add_library(trt_runtime INTERFACE) target_include_directories(trt_runtime INTERFACE "${NVINFER_INCLUDE}" "${CUDART_INCLUDE}") target_link_libraries(trt_runtime INTERFACE - "${NVINFER_LIB}" "${CUDART_LIB}") -message(STATUS "TensorRT: ${NVINFER_LIB} CUDA runtime: ${CUDART_LIB}") + "${NVINFER_LIB}" "${CUDART_LIB}" "${CUBLAS_LIB}") +message(STATUS "TensorRT: ${NVINFER_LIB} CUDA runtime: ${CUDART_LIB} cuBLAS: ${CUBLAS_LIB}") # FFmpeg (NVDEC hardware video decode + swscale colour conversion) find_package(PkgConfig REQUIRED) diff --git a/README.md b/README.md index b05dfed..0b8c6f3 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Identifies actors in movie files and produces X-ray-style scene annotations comp ## How it works 1. **Build a gallery** — download actor headshots from TMDB/IMDB, embed them with ArcFace (`build_gallery` / `scripts/make_gallery.py`). -2. **Analyze a movie** — `scene_analyze` decodes frames at configurable FPS, detects faces (YuNet/SCRFD), tracks them across cuts, matches identities against the gallery using calibrated similarity, and writes time-window JSON. +2. **Analyze a movie** — `scene_analyze` decodes frames at configurable FPS, detects faces (SCRFD), tracks them across cuts, matches identities against the gallery using calibrated similarity, and writes time-window JSON. 3. **Output** — minimal mode produces Jellyfin-ready actor name + time-window JSON; standard mode adds per-frame bbox, similarity, and track data. ## Dependencies @@ -13,9 +13,12 @@ Identifies actors in movie files and produces X-ray-style scene annotations comp | Dependency | Role | |---|---| | KPN++ | Pipeline backbone (nodes, networks) | -| OpenCV 4 | Video decode, image ops, DNN inference, YuNet face detection | +| OpenCV 4 | Video decode, image ops, DNN inference | | ONNX Runtime | SCRFD face detector (dynamic shape nodes unsupported by cv::dnn) | +| TensorRT + CUDA runtime + cuBLAS | Optional TRT engines for SCRFD/ArcFace (`--detector-engine`/`--arcface-engine`); identity_matcher's GPU gallery scan | +| FFmpeg (libav*) | NVDEC hardware video decode + colour conversion | | nlohmann/json | JSON I/O | +| nanobind | Python bindings for `sae_embed` | ## Build @@ -62,6 +65,7 @@ Models are placed in `external/`: | `scene_analyze_debug` | Same as above + per-frame annotated JPEGs (`SAE_DEBUG=1`) | | `scene_preview` | Live OpenCV display window while analysing | | `build_gallery` | Offline gallery builder from a directory of images | +| `embed_faces` | CLI: image(s) → embedding JSON, used by gallery-builder scripts | | `sae_embed` | Python module (nanobind) used by gallery-builder scripts — loads SCRFD+ArcFace once | ### `scene_analyze` @@ -82,7 +86,6 @@ Key options: | `--track-min-iou` | — | Minimum IoU gate for spatial assignment | | `--track-max-embed` | — | Maximum embedding distance gate | | `--track-max-missing` | — | Frames a track survives without a detection | -| `--track-min-frames` | 3 | Observations before a track's mean embedding is used for matching | ### Gallery builders diff --git a/src/config.hpp b/src/config.hpp index 34f2129..efe5baa 100644 --- a/src/config.hpp +++ b/src/config.hpp @@ -53,7 +53,6 @@ struct Config { float track_min_iou{0.1f}; // IoU below which spatial link alone is rejected float track_max_embed_dist{0.7f}; // cosine dist above which embedding link alone is rejected int track_max_frames_missing{5}; // expire track after N consecutive missed frames - int track_min_frames{3}; // frames before track mean replaces per-frame embedding // ── Scene tracking ──────────────────────────────────────────────────────── double extinction_sec{5.0}; // keep actor active this many seconds after last detection diff --git a/src/main.cpp b/src/main.cpp index 65e23fe..4af2ec9 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -94,7 +94,6 @@ 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("--track-min-frames")) cfg.track_min_frames = std::stoi(next()); else if (arg("--anneal")) cfg.anneal_sec = std::stod(next()); else if (arg("--trt-cache")) cfg.trt.cache_dir = next(); else if (arg("--trt-fp16")) cfg.trt.fp16 = true; diff --git a/src/nodes/face_tracker_node.hpp b/src/nodes/face_tracker_node.hpp index ac478cf..61f5c02 100644 --- a/src/nodes/face_tracker_node.hpp +++ b/src/nodes/face_tracker_node.hpp @@ -14,10 +14,8 @@ // algorithm on a combined spatial (IoU) + embedding (cosine distance) cost. // // Each track accumulates a running directional mean of its ArcFace embeddings -// (averaged then re-normalised to the unit sphere). Once a track reaches -// min_frames observations its mean embedding is forwarded as track_embeddings[i] -// and track_mature[i] is set, allowing the identity matcher to use a cleaner, -// multi-frame signal instead of the noisy single-frame embedding. +// (averaged then re-normalised to the unit sphere), used as the embedding side +// of the assignment cost below for more stable track continuity. // // Assignment cost (track i, detection j): // cost = alpha * (1 - IoU) + (1-alpha) * min(cosine_dist/2, 1) @@ -41,13 +39,11 @@ struct FaceTrackerFunc { , min_iou_(cfg.track_min_iou) , max_embed_dist_(cfg.track_max_embed_dist) , max_missing_(cfg.track_max_frames_missing) - , min_frames_(cfg.track_min_frames) { std::cerr << "[face_tracker] alpha=" << alpha_ << " min_iou=" << min_iou_ << " max_embed_dist=" << max_embed_dist_ - << " max_missing=" << max_missing_ - << " min_frames=" << min_frames_ << "\n"; + << " max_missing=" << max_missing_ << "\n"; } TrackedSceneFrame operator()(EmbeddedSceneFrame ef) { @@ -102,8 +98,6 @@ struct FaceTrackerFunc { out.crops = ef.crops; out.embeddings = ef.embeddings; out.track_ids.assign(n_det, -1); - out.track_embeddings = ef.embeddings; // default: per-frame embedding - out.track_mature.assign(n_det, false); std::vector det_matched(n_det, false); @@ -122,9 +116,7 @@ struct FaceTrackerFunc { ts.frames_missing = 0; det_matched[di] = true; - out.track_ids[di] = tids[ti]; - out.track_embeddings[di] = ts.mean_emb; - out.track_mature[di] = (ts.n_frames >= min_frames_); + out.track_ids[di] = tids[ti]; } // Create new tracks for unmatched detections @@ -137,7 +129,6 @@ struct FaceTrackerFunc { ts.n_frames = 1; tracks_[tid] = ts; out.track_ids[di] = tid; - // track_embeddings[di] already initialised to per-frame embedding } // Expire stale tracks @@ -239,5 +230,4 @@ private: float min_iou_; float max_embed_dist_; int max_missing_; - int min_frames_; }; diff --git a/src/nodes/identity_matcher_node.hpp b/src/nodes/identity_matcher_node.hpp index ddd337f..2032e86 100644 --- a/src/nodes/identity_matcher_node.hpp +++ b/src/nodes/identity_matcher_node.hpp @@ -4,9 +4,15 @@ #include "gallery/gallery_store.hpp" #include "gallery/gallery_calibration.hpp" -#include -#include +#include +#include + +#include +#include #include +#include +#include +#include // ── IdentityMatcherFunc ─────────────────────────────────────────────────────── // KPN node: compares each embedding against every reference embedding in the @@ -23,12 +29,38 @@ // (b) ratio test: best_dist/second_best_dist < match_ratio // AND best_dist < match_ratio_ceil. // -// In both modes, per-actor best similarity is determined by scanning all +// In both modes, per-actor best similarity is determined by scanning // reference embeddings and taking the closest (best-of-N). +// +// Gallery scan: the full reference set (tens of thousands of 512-dim +// embeddings) is uploaded to the GPU once at construction time and stays +// resident there. Per frame, only the small query matrix (n_faces x 512) is +// uploaded and a single cublasSgemm computes the full similarity matrix +// (n_faces x N_gallery) in well under a millisecond — far faster than any +// CPU GEMM or scalar scan, making gallery-side pruning unnecessary. + +namespace identity_matcher_detail { +struct CudaError : std::runtime_error { + using std::runtime_error::runtime_error; +}; + +inline void check_cuda(cudaError_t e, const char* what) { + if (e != cudaSuccess) + throw CudaError(std::string(what) + ": " + cudaGetErrorString(e)); +} + +inline void check_cublas(cublasStatus_t s, const char* what) { + if (s != CUBLAS_STATUS_SUCCESS) + throw CudaError(std::string(what) + ": cublas error " + std::to_string(s)); +} +} // namespace identity_matcher_detail struct IdentityMatcherFunc { static constexpr std::string_view label() { return "identity_matcher"; } + // Max faces handled per frame without reallocating GPU buffers. + static constexpr int kMaxFaces = 32; + IdentityMatcherFunc(const ActorGallery& gallery, const Config& cfg) : gallery_(gallery) , prob_threshold_(cfg.prob_threshold) @@ -37,6 +69,8 @@ struct IdentityMatcherFunc { , ratio_(cfg.match_ratio) , ratio_ceil_(cfg.match_ratio_ceil) { + using namespace identity_matcher_detail; + std::cerr << "[identity_matcher] flattening gallery embeddings...\n"; for (int ai = 0; ai < static_cast(gallery_.actors.size()); ++ai) { for (const auto& emb : gallery_.actors[ai].embeddings) { @@ -44,6 +78,7 @@ struct IdentityMatcherFunc { flat_actor_.push_back(ai); } } + n_gallery_ = static_cast(flat_emb_.size()); std::cerr << "[identity_matcher] starting calibration (" << flat_emb_.size() << " embeddings)...\n"; @@ -64,27 +99,93 @@ struct IdentityMatcherFunc { std::cerr << "[identity_matcher] gallery: " << gallery_.actors.size() << " actors, " << flat_emb_.size() << " reference embeddings\n"; + + // Flatten gallery into a contiguous (N x 512) row-major host buffer, + // then upload once. Row-major NxD == column-major DxN, which is the + // layout cublasSgemm wants for the transposed operand below. + std::vector host_gallery(static_cast(n_gallery_) * 512); + for (int i = 0; i < n_gallery_; ++i) + std::memcpy(host_gallery.data() + static_cast(i) * 512, + flat_emb_[i].data(), 512 * sizeof(float)); + + check_cuda(cudaMalloc(reinterpret_cast(&d_gallery_), host_gallery.size() * sizeof(float)), + "cudaMalloc gallery"); + check_cuda(cudaMemcpy(d_gallery_, host_gallery.data(), + host_gallery.size() * sizeof(float), + cudaMemcpyHostToDevice), + "cudaMemcpy gallery H2D"); + + check_cuda(cudaMalloc(reinterpret_cast(&d_query_), static_cast(kMaxFaces) * 512 * sizeof(float)), + "cudaMalloc query"); + check_cuda(cudaMalloc(reinterpret_cast(&d_sims_), static_cast(kMaxFaces) * n_gallery_ * sizeof(float)), + "cudaMalloc sims"); + + check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate"); + check_cublas(cublasCreate(&handle_), "cublasCreate"); + check_cublas(cublasSetStream(handle_, stream_), "cublasSetStream"); + + host_sims_.resize(static_cast(kMaxFaces) * n_gallery_); + + std::cerr << "[identity_matcher] gallery resident on GPU (" + << (host_gallery.size() * sizeof(float)) / (1024 * 1024) << " MiB)\n"; + } + + ~IdentityMatcherFunc() { + if (d_gallery_) cudaFree(d_gallery_); + if (d_query_) cudaFree(d_query_); + if (d_sims_) cudaFree(d_sims_); + if (handle_) cublasDestroy(handle_); + if (stream_) cudaStreamDestroy(stream_); } MatchedSceneFrame operator()(TrackedSceneFrame tf) { + using namespace identity_matcher_detail; + if (tf.source.eof) return {std::move(tf.source), {}}; + const int n_faces = static_cast(tf.embeddings.size()); std::vector actors; - actors.reserve(tf.embeddings.size()); + actors.reserve(n_faces); - for (int fi = 0; fi < static_cast(tf.embeddings.size()); ++fi) { - // Prefer the track's accumulated mean embedding when the track is - // mature (≥ min_frames observations) — more stable than single-frame. - const Embedding& query = tf.track_mature[fi] - ? tf.track_embeddings[fi] - : tf.embeddings[fi]; + if (n_faces == 0) return {std::move(tf.source), {}}; + if (n_faces > kMaxFaces) + throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces"); - // Per-actor best cosine similarity (max dot product) + // Build the (n_faces x 512) query matrix (row-major == col-major 512 x n_faces). + std::vector host_query(static_cast(n_faces) * 512); + for (int fi = 0; fi < n_faces; ++fi) { + std::memcpy(host_query.data() + static_cast(fi) * 512, + tf.embeddings[fi].data(), 512 * sizeof(float)); + } + + check_cuda(cudaMemcpyAsync(d_query_, host_query.data(), + host_query.size() * sizeof(float), + cudaMemcpyHostToDevice, stream_), + "cudaMemcpy query H2D"); + + // S (N_gallery x n_faces) col-major = G(512 x N_gallery)^T * Q(512 x n_faces) + // i.e. S[g + f*N_gallery] = cosine_similarity(gallery[g], query[f]). + const float alpha = 1.f, beta = 0.f; + check_cublas(cublasSgemm(handle_, CUBLAS_OP_T, CUBLAS_OP_N, + n_gallery_, n_faces, 512, + &alpha, d_gallery_, 512, d_query_, 512, + &beta, d_sims_, n_gallery_), + "cublasSgemm"); + + check_cuda(cudaMemcpyAsync(host_sims_.data(), d_sims_, + static_cast(n_gallery_) * n_faces * sizeof(float), + cudaMemcpyDeviceToHost, stream_), + "cudaMemcpy sims D2H"); + check_cuda(cudaStreamSynchronize(stream_), "cudaStreamSynchronize"); + + for (int fi = 0; fi < n_faces; ++fi) { + const float* sims = host_sims_.data() + static_cast(fi) * n_gallery_; + + // Per-actor best cosine similarity (max over that actor's reference embeddings) std::vector best_sim(gallery_.actors.size(), -std::numeric_limits::max()); - - for (int ei = 0; ei < static_cast(flat_emb_.size()); ++ei) { - float sim = cosine_similarity(query, flat_emb_[ei]); + for (int ei = 0; ei < n_gallery_; ++ei) { + float sim = sims[ei]; int ai = flat_actor_[ei]; if (sim > best_sim[ai]) best_sim[ai] = sim; } @@ -157,4 +258,13 @@ private: float ratio_ceil_; std::vector flat_emb_; std::vector flat_actor_; + int n_gallery_{0}; + + float* d_gallery_{nullptr}; // (n_gallery_ x 512), resident for the lifetime of this node + float* d_query_{nullptr}; // (kMaxFaces x 512) + float* d_sims_{nullptr}; // (n_gallery_ x kMaxFaces), column-major + std::vector host_sims_; + + cudaStream_t stream_{nullptr}; + cublasHandle_t handle_{nullptr}; }; diff --git a/src/scene_preview.cpp b/src/scene_preview.cpp index 7bce18e..8fd4236 100644 --- a/src/scene_preview.cpp +++ b/src/scene_preview.cpp @@ -79,7 +79,6 @@ 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("--track-min-frames")) cfg.track_min_frames = std::stoi(next()); else if (arg("--anneal")) cfg.anneal_sec = std::stod(next()); else if (arg("--trt-cache")) cfg.trt.cache_dir = next(); else if (arg("--trt-fp16")) cfg.trt.fp16 = true; diff --git a/src/types.hpp b/src/types.hpp index 73c26e5..38c13ed 100644 --- a/src/types.hpp +++ b/src/types.hpp @@ -72,8 +72,7 @@ struct EmbeddedSceneFrame { // ── Face tracking ───────────────────────────────────────────────────────────── // Output of FaceTrackerFunc — EmbeddedSceneFrame augmented with per-detection -// track context. track_embeddings[i] is the L2-normalised running mean across -// the track's history; use it for identity matching when track_mature[i] is true. +// track context. struct TrackedSceneFrame { Frame source; @@ -81,8 +80,6 @@ struct TrackedSceneFrame { std::vector crops; std::vector track_ids; // -1 = brand-new track this frame std::vector embeddings; // per-frame raw (from embedder) - std::vector track_embeddings; // accumulated mean per track - std::vector track_mature; // true once track has ≥ min_frames obs. }; // ── Identity matching ─────────────────────────────────────────────────────────