improved performance
This commit is contained in:
+7
-4
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<bool> det_matched(n_det, false);
|
||||
|
||||
@@ -123,8 +117,6 @@ struct FaceTrackerFunc {
|
||||
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_);
|
||||
}
|
||||
|
||||
// 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_;
|
||||
};
|
||||
|
||||
@@ -4,9 +4,15 @@
|
||||
#include "gallery/gallery_store.hpp"
|
||||
#include "gallery/gallery_calibration.hpp"
|
||||
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <cublas_v2.h>
|
||||
#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) {
|
||||
// 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<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 < static_cast<int>(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<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};
|
||||
};
|
||||
|
||||
@@ -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;
|
||||
|
||||
+1
-4
@@ -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<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 ─────────────────────────────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user