Initial commit: scene-actor-extraction pipeline
Source (KPN++ pipeline nodes, ArcFace embedders, SCRFD/YuNet detectors, gallery builder), build scripts, and eval artifacts. - external/KPN as a git submodule (gitea.tourolle.paris/dtourolle/KPN) - ONNX models tracked via Git LFS (models/*.onnx) - generated outputs, TensorRT engines, reference repos, and media ignored
This commit is contained in:
@@ -0,0 +1,243 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <limits>
|
||||
#include <map>
|
||||
#include <vector>
|
||||
|
||||
// ── FaceTrackerFunc ───────────────────────────────────────────────────────────
|
||||
// KPN node: links face detections across consecutive frames using the Hungarian
|
||||
// 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.
|
||||
//
|
||||
// Assignment cost (track i, detection j):
|
||||
// cost = alpha * (1 - IoU) + (1-alpha) * min(cosine_dist/2, 1)
|
||||
// Gated to INF when IoU < min_iou AND cosine_dist > max_embed_dist.
|
||||
//
|
||||
// Unmatched tracks have their frames_missing counter incremented; they are
|
||||
// expired once frames_missing > max_frames_missing.
|
||||
|
||||
struct FaceTrackerFunc {
|
||||
static constexpr std::string_view label() { return "face_tracker"; }
|
||||
|
||||
struct TrackState {
|
||||
cv::Rect2f bbox;
|
||||
Embedding mean_emb{};
|
||||
int n_frames{0};
|
||||
int frames_missing{0};
|
||||
};
|
||||
|
||||
explicit FaceTrackerFunc(const Config& cfg)
|
||||
: alpha_(cfg.track_alpha)
|
||||
, 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";
|
||||
}
|
||||
|
||||
TrackedSceneFrame operator()(EmbeddedSceneFrame ef) {
|
||||
if (ef.source.eof) {
|
||||
tracks_.clear();
|
||||
TrackedSceneFrame out;
|
||||
out.source = std::move(ef.source);
|
||||
return out;
|
||||
}
|
||||
|
||||
const int n_det = static_cast<int>(ef.embeddings.size());
|
||||
|
||||
if (ef.source.is_cut && !tracks_.empty()) {
|
||||
std::cerr << "[face_tracker] cut — clearing " << tracks_.size() << " tracks\n";
|
||||
tracks_.clear();
|
||||
}
|
||||
|
||||
// Snapshot active track IDs so the map can be modified safely below
|
||||
std::vector<int> tids;
|
||||
tids.reserve(tracks_.size());
|
||||
for (auto& [tid, _] : tracks_) tids.push_back(tid);
|
||||
const int n_trk = static_cast<int>(tids.size());
|
||||
|
||||
// ── Cost matrix [n_trk × n_det] ──────────────────────────────────────
|
||||
constexpr float INF_COST = 1e6f;
|
||||
std::vector<std::vector<float>> cost(n_trk,
|
||||
std::vector<float>(n_det, INF_COST));
|
||||
|
||||
for (int ti = 0; ti < n_trk; ++ti) {
|
||||
const TrackState& ts = tracks_[tids[ti]];
|
||||
for (int di = 0; di < n_det; ++di) {
|
||||
float iou_v = iou(ts.bbox, ef.faces[di].bbox);
|
||||
float emb_d = (ts.n_frames > 0)
|
||||
? 1.f - cosine_similarity(ts.mean_emb, ef.embeddings[di])
|
||||
: 1.f;
|
||||
if (iou_v < min_iou_ && emb_d > max_embed_dist_) continue;
|
||||
float s = 1.f - iou_v;
|
||||
float e = std::min(emb_d * 0.5f, 1.f);
|
||||
cost[ti][di] = alpha_ * s + (1.f - alpha_) * e;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Hungarian assignment ──────────────────────────────────────────────
|
||||
std::vector<int> assign(n_trk, -1);
|
||||
if (n_trk > 0 && n_det > 0)
|
||||
assign = hungarian(cost, n_trk, n_det);
|
||||
|
||||
// ── Build output frame ────────────────────────────────────────────────
|
||||
TrackedSceneFrame out;
|
||||
out.source = ef.source;
|
||||
out.faces = ef.faces;
|
||||
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);
|
||||
|
||||
// Update matched tracks
|
||||
for (int ti = 0; ti < n_trk; ++ti) {
|
||||
int di = assign[ti];
|
||||
bool valid = (di >= 0 && di < n_det && cost[ti][di] < INF_COST * 0.5f);
|
||||
TrackState& ts = tracks_[tids[ti]];
|
||||
if (!valid) {
|
||||
ts.frames_missing++;
|
||||
continue;
|
||||
}
|
||||
update_mean(ts.mean_emb, ts.n_frames, ef.embeddings[di]);
|
||||
ts.bbox = ef.faces[di].bbox;
|
||||
ts.n_frames++;
|
||||
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_);
|
||||
}
|
||||
|
||||
// Create new tracks for unmatched detections
|
||||
for (int di = 0; di < n_det; ++di) {
|
||||
if (det_matched[di]) continue;
|
||||
int tid = next_id_++;
|
||||
TrackState ts;
|
||||
ts.bbox = ef.faces[di].bbox;
|
||||
ts.mean_emb = ef.embeddings[di];
|
||||
ts.n_frames = 1;
|
||||
tracks_[tid] = ts;
|
||||
out.track_ids[di] = tid;
|
||||
// track_embeddings[di] already initialised to per-frame embedding
|
||||
}
|
||||
|
||||
// Expire stale tracks
|
||||
for (auto it = tracks_.begin(); it != tracks_.end(); ) {
|
||||
it = (it->second.frames_missing > max_missing_)
|
||||
? tracks_.erase(it) : std::next(it);
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
private:
|
||||
// IoU of two axis-aligned bounding boxes
|
||||
static float iou(const cv::Rect2f& a, const cv::Rect2f& b) {
|
||||
float ix = std::max(0.f, std::min(a.x + a.width, b.x + b.width)
|
||||
- std::max(a.x, b.x));
|
||||
float iy = std::max(0.f, std::min(a.y + a.height, b.y + b.height)
|
||||
- std::max(a.y, b.y));
|
||||
float inter = ix * iy;
|
||||
if (inter <= 0.f) return 0.f;
|
||||
return inter / (a.width * a.height + b.width * b.height - inter);
|
||||
}
|
||||
|
||||
// Online directional mean: average then re-normalise to unit sphere
|
||||
static void update_mean(Embedding& mean, int n_prev, const Embedding& emb) {
|
||||
float norm_sq = 0.f;
|
||||
for (int k = 0; k < 512; ++k) {
|
||||
mean[k] = (mean[k] * n_prev + emb[k]) / (n_prev + 1);
|
||||
norm_sq += mean[k] * mean[k];
|
||||
}
|
||||
float inv = 1.f / std::sqrt(norm_sq);
|
||||
for (int k = 0; k < 512; ++k) mean[k] *= inv;
|
||||
}
|
||||
|
||||
// O(n³) potential-based Hungarian algorithm (Jonker-Volgenant / Kuhn-Munkres).
|
||||
// Returns assign[row] = col (0-indexed), or -1 when row is matched to a
|
||||
// padded virtual column (i.e., unmatched). Rectangular matrices are padded
|
||||
// to square with 0-cost virtual entries so leftover rows/cols are absorbed
|
||||
// cheaply rather than being forced onto real rows/cols.
|
||||
static std::vector<int> hungarian(
|
||||
const std::vector<std::vector<float>>& C, int nr, int nc)
|
||||
{
|
||||
const int N = std::max(nr, nc);
|
||||
constexpr float INF_VAL = 1e30f;
|
||||
|
||||
// Expand to N×N, filling virtual entries with 0
|
||||
std::vector<std::vector<float>> sq(N, std::vector<float>(N, 0.f));
|
||||
for (int i = 0; i < nr; ++i)
|
||||
for (int j = 0; j < nc; ++j)
|
||||
sq[i][j] = C[i][j];
|
||||
|
||||
std::vector<float> u(N + 1, 0.f), v(N + 1, 0.f);
|
||||
std::vector<int> p(N + 1, 0), way(N + 1, 0);
|
||||
|
||||
for (int i = 1; i <= N; ++i) {
|
||||
p[0] = i;
|
||||
int j0 = 0;
|
||||
std::vector<float> minv(N + 1, INF_VAL);
|
||||
std::vector<bool> used(N + 1, false);
|
||||
do {
|
||||
used[j0] = true;
|
||||
int i0 = p[j0], j1 = -1;
|
||||
float delta = INF_VAL;
|
||||
for (int j = 1; j <= N; ++j) {
|
||||
if (!used[j]) {
|
||||
float cur = sq[i0-1][j-1] - u[i0] - v[j];
|
||||
if (cur < minv[j]) { minv[j] = cur; way[j] = j0; }
|
||||
if (minv[j] < delta) { delta = minv[j]; j1 = j; }
|
||||
}
|
||||
}
|
||||
for (int j = 0; j <= N; ++j) {
|
||||
if (used[j]) { u[p[j]] += delta; v[j] -= delta; }
|
||||
else minv[j] -= delta;
|
||||
}
|
||||
j0 = j1;
|
||||
} while (p[j0] != 0);
|
||||
do {
|
||||
int j1 = way[j0];
|
||||
p[j0] = p[j1];
|
||||
j0 = j1;
|
||||
} while (j0);
|
||||
}
|
||||
|
||||
// p[j] = row (1-indexed) assigned to column j (1-indexed)
|
||||
std::vector<int> ans(nr, -1);
|
||||
for (int j = 1; j <= N; ++j) {
|
||||
int row = p[j] - 1;
|
||||
int col = j - 1;
|
||||
if (row >= 0 && row < nr && col < nc)
|
||||
ans[row] = col;
|
||||
// col >= nc → virtual column → row stays unmatched (-1)
|
||||
}
|
||||
return ans;
|
||||
}
|
||||
|
||||
std::map<int, TrackState> tracks_;
|
||||
int next_id_{0};
|
||||
float alpha_;
|
||||
float min_iou_;
|
||||
float max_embed_dist_;
|
||||
int max_missing_;
|
||||
int min_frames_;
|
||||
};
|
||||
Reference in New Issue
Block a user