Closes the link that made AR-012 inert: the tracker was maintaining registry state, but nothing called observe(), so no belief accumulated, no track was ever owned, and no presence claim could be emitted. Tracking worked and presence did not. The matcher now feeds every scored face to the registry as a calibrated posterior plus its embedding. Deliberately every scored face, not only the ones clearing prob_threshold: a run of near-misses for one actor is evidence, and discarding it would leave ownership depending on the per-frame threshold this redesign exists to stop relying on. The registry discounts for correlation and decides ownership from the accumulated posterior (AR-025). The registry is an optional dependency of the matcher. Without one it behaves exactly as before, which keeps the replay harness and the unit tests working unchanged rather than forcing every caller to construct a registry it does not need. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> TRACES: AR-012, AR-025 | SR-002
280 lines
13 KiB
C++
280 lines
13 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 "gallery/track_gallery.hpp"
|
||
#include "track_registry.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)
|
||
, track_gallery_(cfg)
|
||
{
|
||
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";
|
||
bool recomputed = false;
|
||
cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_,
|
||
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"
|
||
<< " 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);
|
||
}
|
||
|
||
/// TRACES: AR-023, AR-024 | SR-002
|
||
/// The fitted sigmoid. Exposed because the matcher is where it gets fitted
|
||
/// (and cached back to the gallery), but it is not the matcher's private
|
||
/// property: track association and evidence weighting must threshold in the
|
||
/// *same* probability space, or a "0.5" in one stage and a "0.5" in another
|
||
/// mean different things. See `same_person_probability`.
|
||
const GalleryCalibration& calibration() const { return cal_; }
|
||
|
||
/// TRACES: AR-012, AR-025 | SR-002
|
||
/// Where per-frame identity evidence reaches the registry. Optional: with no
|
||
/// registry attached the matcher behaves exactly as before, which keeps the
|
||
/// replay harness and the unit tests working unchanged.
|
||
void set_registry(std::shared_ptr<TrackRegistry> r) { registry_ = std::move(r); }
|
||
|
||
// 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) {
|
||
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<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;
|
||
}
|
||
|
||
// 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<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;
|
||
// 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];
|
||
|
||
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;
|
||
}
|
||
|
||
// 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.
|
||
// TRACES: AR-012, AR-025 | SR-002
|
||
// Every scored face is evidence, not only the accepted ones: a run of
|
||
// near-misses for one actor is itself informative, and discarding it
|
||
// would make ownership depend on a per-frame threshold the redesign
|
||
// exists to stop relying on. The registry discounts for correlation
|
||
// and decides ownership from the accumulated posterior (AR-025).
|
||
if (registry_ && best_actor >= 0 && tf.track_ids[fi] >= 0) {
|
||
const float p = cal_.valid
|
||
? cal_.probability(best_s, log_prior_odds_)
|
||
: std::max(0.f, best_s);
|
||
registry_->observe(tf.track_ids[fi], best_actor, p, tf.embeddings[fi]);
|
||
}
|
||
|
||
track_gallery_.observe(tf.track_ids[fi], tf.embeddings[fi],
|
||
best_actor, best_s, accept, tf.crops[fi]);
|
||
|
||
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_;
|
||
TrackGallery track_gallery_;
|
||
std::shared_ptr<TrackRegistry> registry_;
|
||
};
|