193 lines
8.1 KiB
C++
193 lines
8.1 KiB
C++
#pragma once
|
||
#include "types.hpp"
|
||
#include "config.hpp"
|
||
#include "inference/similarity.hpp"
|
||
#include "gallery/gallery_store.hpp"
|
||
#include "gallery/gallery_calibration.hpp"
|
||
|
||
#include <cstdint>
|
||
#include <cstring>
|
||
#include <iostream>
|
||
#include <limits>
|
||
#include <memory>
|
||
#include <stdexcept>
|
||
#include <vector>
|
||
|
||
// ── IdentityMatcherFunc ───────────────────────────────────────────────────────
|
||
// KPN node: compares each embedding against every reference embedding in the
|
||
// actor gallery using cosine similarity.
|
||
//
|
||
// Matching strategy — two modes selected at construction time:
|
||
//
|
||
// Calibrated (preferred): gallery calibration fits a sigmoid
|
||
// P(match) = σ(a·similarity + b) from intra/inter-class pairs.
|
||
// A face is accepted if P(match | best_actor) > prob_threshold.
|
||
//
|
||
// Fallback (no calibration): dual-criterion accept —
|
||
// (a) best cosine distance < match_threshold, OR
|
||
// (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
|
||
// 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 SGEMM computes the full similarity matrix in well under
|
||
// a millisecond. The GPU math backend (cuBLAS or rocBLAS) lives behind
|
||
// ISimilarityEngine (backends/gemm_backend.cpp) and is selected at compile time.
|
||
|
||
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)
|
||
, log_prior_odds_(std::log(cfg.match_prior / (1.f - cfg.match_prior)))
|
||
, threshold_(cfg.match_threshold)
|
||
, ratio_(cfg.match_ratio)
|
||
, ratio_ceil_(cfg.match_ratio_ceil)
|
||
{
|
||
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) {
|
||
flat_emb_.push_back(emb);
|
||
flat_actor_.push_back(ai);
|
||
}
|
||
}
|
||
n_gallery_ = static_cast<int>(flat_emb_.size());
|
||
|
||
std::cerr << "[identity_matcher] starting calibration ("
|
||
<< flat_emb_.size() << " embeddings)...\n";
|
||
cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_,
|
||
cfg.gallery_path + ".calib_cache.json");
|
||
|
||
if (cal_.valid) {
|
||
std::cerr << "[identity_matcher] calibrated Bayesian matching"
|
||
<< " prior=" << cfg.match_prior
|
||
<< " P_threshold=" << prob_threshold_
|
||
<< " effective_sim_boundary="
|
||
<< cal_.boundary_at(prob_threshold_, log_prior_odds_) << "\n";
|
||
} else {
|
||
std::cerr << "[identity_matcher] threshold matching (calibration skipped)"
|
||
<< " threshold=" << threshold_
|
||
<< " ratio=" << ratio_ << " ratio_ceil=" << ratio_ceil_ << "\n";
|
||
}
|
||
std::cerr << "[identity_matcher] gallery: "
|
||
<< gallery_.actors.size() << " actors, "
|
||
<< flat_emb_.size() << " reference embeddings\n";
|
||
|
||
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));
|
||
|
||
sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces);
|
||
}
|
||
|
||
MatchedSceneFrame operator()(TrackedSceneFrame tf) {
|
||
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(n_faces);
|
||
|
||
if (n_faces == 0) return {std::move(tf.source), {}};
|
||
if (n_faces > kMaxFaces)
|
||
throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");
|
||
|
||
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));
|
||
}
|
||
|
||
// S (N_gallery × n_faces) col-major: face fi's gallery sims at sims + fi*n_gallery.
|
||
const float* host_sims = sim_engine_->compute(host_query.data(), n_faces);
|
||
|
||
for (int fi = 0; fi < n_faces; ++fi) {
|
||
const float* sims = host_sims + static_cast<size_t>(fi) * n_gallery_;
|
||
|
||
std::vector<float> best_sim(gallery_.actors.size(),
|
||
-std::numeric_limits<float>::max());
|
||
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;
|
||
}
|
||
|
||
int best_actor = -1;
|
||
int second_actor = -1;
|
||
float best_s = -std::numeric_limits<float>::max();
|
||
float second_s = -std::numeric_limits<float>::max();
|
||
for (int ai = 0; ai < static_cast<int>(best_sim.size()); ++ai) {
|
||
if (best_sim[ai] > best_s) {
|
||
second_s = best_s;
|
||
second_actor = best_actor;
|
||
best_s = best_sim[ai];
|
||
best_actor = ai;
|
||
} else if (best_sim[ai] > second_s) {
|
||
second_s = best_sim[ai];
|
||
second_actor = ai;
|
||
}
|
||
}
|
||
(void)second_actor;
|
||
|
||
bool accept = false;
|
||
if (best_actor >= 0) {
|
||
if (cal_.valid) {
|
||
accept = cal_.probability(best_s, log_prior_odds_) > prob_threshold_;
|
||
} else {
|
||
float best_d = 1.f - best_s;
|
||
float second_d = (second_s > -std::numeric_limits<float>::max())
|
||
? 1.f - second_s
|
||
: std::numeric_limits<float>::max();
|
||
bool absolute = best_d < threshold_;
|
||
bool ratio = (best_d < ratio_ceil_) &&
|
||
(second_d == std::numeric_limits<float>::max() ||
|
||
best_d / second_d < ratio_);
|
||
accept = absolute || ratio;
|
||
}
|
||
}
|
||
|
||
IdentifiedActor ia;
|
||
ia.bbox = tf.faces[fi].bbox;
|
||
ia.crop = tf.crops[fi];
|
||
ia.track_id = tf.track_ids[fi];
|
||
|
||
if (accept) {
|
||
ia.actor_idx = best_actor;
|
||
ia.name = gallery_.actors[best_actor].name;
|
||
ia.imdb_id = gallery_.actors[best_actor].imdb_id;
|
||
ia.tmdb_id = gallery_.actors[best_actor].tmdb_id;
|
||
ia.jellyfin_id = gallery_.actors[best_actor].jellyfin_id;
|
||
ia.similarity = cal_.valid
|
||
? cal_.probability(best_s, log_prior_odds_)
|
||
: best_s;
|
||
}
|
||
|
||
actors.push_back(std::move(ia));
|
||
}
|
||
|
||
return {std::move(tf.source), std::move(actors)};
|
||
}
|
||
|
||
private:
|
||
ActorGallery gallery_;
|
||
GalleryCalibration cal_;
|
||
float prob_threshold_;
|
||
float log_prior_odds_;
|
||
float threshold_;
|
||
float ratio_;
|
||
float ratio_ceil_;
|
||
std::vector<Embedding> flat_emb_;
|
||
std::vector<int> flat_actor_;
|
||
int n_gallery_{0};
|
||
|
||
std::unique_ptr<ISimilarityEngine> sim_engine_;
|
||
};
|