improved performance

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}")
+    "${NVINFER_LIB}" "${CUDART_LIB}" "${CUBLAS_LIB}")
-message(STATUS "TensorRT: ${NVINFER_LIB}  CUDA runtime: ${CUDART_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)

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
 

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

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;

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
+// (averaged then re-normalised to the unit sphere), used as the embedding side
-// min_frames observations its mean embedding is forwarded as track_embeddings[i]
+// of the assignment cost below for more stable track continuity.
-// and track_mature[i] is set, allowing the identity matcher to use a cleaner,
-// multi-frame signal instead of the noisy single-frame embedding.
 //
 // 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_
+                  << "  max_missing=" << max_missing_ << "\n";
-                  << "  min_frames=" << min_frames_ << "\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<bool> 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_ids[di] = tids[ti];
-            out.track_embeddings[di] = ts.mean_emb;
-            out.track_mature[di]     = (ts.n_frames >= min_frames_);
         }
 
         // 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_;
 };

src/nodes/identity_matcher_node.hpp

@@ -4,9 +4,15 @@
 #include "gallery/gallery_store.hpp"
 #include "gallery/gallery_calibration.hpp"
 
-#include <cmath>
+#include <cublas_v2.h>
-#include <limits>
+#include <cuda_runtime_api.h>
+
+#include <cstdint>
+#include <cstring>
 #include <iostream>
+#include <limits>
+#include <stdexcept>
+#include <vector>
 
 // ── 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<int>(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<int>(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<float> host_gallery(static_cast<size_t>(n_gallery_) * 512);
+        for (int i = 0; i < n_gallery_; ++i)
+            std::memcpy(host_gallery.data() + static_cast<size_t>(i) * 512,
+                        flat_emb_[i].data(), 512 * sizeof(float));
+
+        check_cuda(cudaMalloc(reinterpret_cast<void**>(&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<void**>(&d_query_), static_cast<size_t>(kMaxFaces) * 512 * sizeof(float)),
+                   "cudaMalloc query");
+        check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_sims_), static_cast<size_t>(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<size_t>(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<int>(tf.embeddings.size());
         std::vector<IdentifiedActor> actors;
-        actors.reserve(tf.embeddings.size());
+        actors.reserve(n_faces);
 
-        for (int fi = 0; fi < static_cast<int>(tf.embeddings.size()); ++fi) {
+        if (n_faces == 0) return {std::move(tf.source), {}};
-            // Prefer the track's accumulated mean embedding when the track is
+        if (n_faces > kMaxFaces)
-            // mature (≥ min_frames observations) — more stable than single-frame.
+            throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");
-            const Embedding& query = tf.track_mature[fi]
-                                     ? tf.track_embeddings[fi]
-                                     : tf.embeddings[fi];
 
-            // 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<float> host_query(static_cast<size_t>(n_faces) * 512);
+        for (int fi = 0; fi < n_faces; ++fi) {
+            std::memcpy(host_query.data() + static_cast<size_t>(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<size_t>(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<size_t>(fi) * n_gallery_;
+
+            // Per-actor best cosine similarity (max over that actor's reference embeddings)
             std::vector<float> best_sim(gallery_.actors.size(),
                                         -std::numeric_limits<float>::max());
-
+            for (int ei = 0; ei < n_gallery_; ++ei) {
-            for (int ei = 0; ei < static_cast<int>(flat_emb_.size()); ++ei) {
+                float sim = sims[ei];
-                float sim = cosine_similarity(query, flat_emb_[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<Embedding> flat_emb_;
     std::vector<int>       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<float> host_sims_;
+
+    cudaStream_t   stream_{nullptr};
+    cublasHandle_t handle_{nullptr};
 };

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;

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
+// track context.
-// the track's history; use it for identity matching when track_mature[i] is true.
 
 struct TrackedSceneFrame {
     Frame                     source;
@@ -81,8 +80,6 @@ struct TrackedSceneFrame {
     std::vector<cv::Mat>      crops;
     std::vector<int>          track_ids;          // -1 = brand-new track this frame
     std::vector<Embedding>    embeddings;          // per-frame raw (from embedder)
-    std::vector<Embedding>    track_embeddings;   // accumulated mean per track
-    std::vector<bool>         track_mature;        // true once track has ≥ min_frames obs.
 };
 
 // ── Identity matching ─────────────────────────────────────────────────────────