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,150 @@
|
||||
#pragma once
|
||||
#ifdef SAE_DEBUG
|
||||
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
// ── DebugRendererFunc ─────────────────────────────────────────────────────────
|
||||
// KPN sink node (SAE_DEBUG only): saves one directory of debug images per frame.
|
||||
//
|
||||
// Output layout:
|
||||
// debug_frames/
|
||||
// t0001.000/
|
||||
// annotated.jpg original frame + coloured bboxes + name overlays
|
||||
// brad_pitt_0.82.jpg 112×112 aligned crop | 1.5× context crop (side-by-side)
|
||||
// unknown_0_0.77.jpg same for unidentified faces
|
||||
//
|
||||
// The node taps MatchedSceneFrame (before scene tracking) so every raw
|
||||
// detection — including unknowns — is captured here.
|
||||
|
||||
struct DebugRendererFunc {
|
||||
static constexpr std::string_view label() { return "debug_renderer"; }
|
||||
|
||||
explicit DebugRendererFunc(const Config& cfg)
|
||||
: cfg_(cfg)
|
||||
{
|
||||
fs::create_directories(cfg_.debug_dir);
|
||||
std::cerr << "[debug_renderer] output dir: " << cfg_.debug_dir << "\n";
|
||||
}
|
||||
|
||||
void operator()(MatchedSceneFrame mf) {
|
||||
if (mf.source.eof || mf.source.image.empty()) return;
|
||||
|
||||
// Directory for this timestamp, e.g. "debug_frames/t0042.000/"
|
||||
char buf[32];
|
||||
std::snprintf(buf, sizeof(buf), "t%08.3f", mf.source.timestamp_sec);
|
||||
fs::path dir = fs::path(cfg_.debug_dir) / buf;
|
||||
fs::create_directories(dir);
|
||||
|
||||
// ── Annotated frame ───────────────────────────────────────────────────
|
||||
cv::Mat annotated = mf.source.image.clone();
|
||||
int unknown_idx = 0;
|
||||
for (const auto& ia : mf.actors) {
|
||||
bool known = (ia.actor_idx >= 0);
|
||||
cv::Scalar colour = known
|
||||
? cv::Scalar(0, 200, 60) // green for identified
|
||||
: cv::Scalar(0, 100, 220); // orange for unknown
|
||||
|
||||
cv::Rect2f b = ia.bbox;
|
||||
cv::rectangle(annotated, b, colour, 2);
|
||||
|
||||
std::string lbl = known
|
||||
? (ia.name + " " + fmt_pct(ia.similarity))
|
||||
: ("unknown " + fmt_pct(ia.similarity));
|
||||
|
||||
// Background strip for readability
|
||||
int baseline = 0;
|
||||
cv::Size ts = cv::getTextSize(lbl, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);
|
||||
cv::Rect strip(static_cast<int>(b.x), static_cast<int>(b.y) - ts.height - 4,
|
||||
ts.width + 4, ts.height + 6);
|
||||
strip &= cv::Rect(0, 0, annotated.cols, annotated.rows);
|
||||
if (strip.area() > 0)
|
||||
cv::rectangle(annotated, strip, colour, cv::FILLED);
|
||||
|
||||
cv::putText(annotated, lbl,
|
||||
cv::Point(static_cast<int>(b.x) + 2, static_cast<int>(b.y) - 2),
|
||||
cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1,
|
||||
cv::LINE_AA);
|
||||
|
||||
// ── Per-face crop image ───────────────────────────────────────────
|
||||
// Left panel: 112×112 aligned crop. Right panel: expanded context.
|
||||
cv::Mat panel = make_face_panel(mf.source.image, ia);
|
||||
|
||||
std::string stem = known
|
||||
? (sanitise(ia.name) + "_" + fmt_sim(ia.similarity))
|
||||
: ("unknown_" + std::to_string(unknown_idx++) + "_" + fmt_sim(ia.similarity));
|
||||
cv::imwrite((dir / (stem + ".jpg")).string(), panel,
|
||||
{cv::IMWRITE_JPEG_QUALITY, 90});
|
||||
}
|
||||
|
||||
cv::imwrite((dir / "annotated.jpg").string(), annotated,
|
||||
{cv::IMWRITE_JPEG_QUALITY, 90});
|
||||
}
|
||||
|
||||
private:
|
||||
const Config& cfg_;
|
||||
|
||||
// Build a side-by-side panel: [112×112 aligned crop | context crop resized to 112×112]
|
||||
cv::Mat make_face_panel(const cv::Mat& frame, const IdentifiedActor& ia) const {
|
||||
// Left: aligned 112×112
|
||||
cv::Mat left = ia.crop.empty()
|
||||
? cv::Mat(112, 112, CV_8UC3, cv::Scalar(60, 60, 60))
|
||||
: ia.crop.clone();
|
||||
|
||||
// Right: expanded bbox from original frame, resized to 112×112
|
||||
cv::Rect2f expanded = expand_bbox(ia.bbox, cfg_.crop_context,
|
||||
frame.cols, frame.rows);
|
||||
cv::Mat right_raw = frame(expanded).clone();
|
||||
cv::Mat right;
|
||||
cv::resize(right_raw, right, {112, 112}, 0, 0, cv::INTER_LINEAR);
|
||||
|
||||
// Separator line
|
||||
cv::Mat sep(112, 4, CV_8UC3, cv::Scalar(200, 200, 200));
|
||||
|
||||
cv::Mat panel;
|
||||
cv::hconcat(std::vector<cv::Mat>{left, sep, right}, panel);
|
||||
return panel;
|
||||
}
|
||||
|
||||
static cv::Rect2f expand_bbox(cv::Rect2f b, float factor, int W, int H) {
|
||||
float cx = b.x + b.width * 0.5f;
|
||||
float cy = b.y + b.height * 0.5f;
|
||||
float nw = b.width * factor;
|
||||
float nh = b.height * factor;
|
||||
float x = std::max(0.f, cx - nw * 0.5f);
|
||||
float y = std::max(0.f, cy - nh * 0.5f);
|
||||
nw = std::min(nw, (float)W - x);
|
||||
nh = std::min(nh, (float)H - y);
|
||||
return {x, y, nw, nh};
|
||||
}
|
||||
|
||||
static std::string fmt_pct(float v) {
|
||||
char buf[8];
|
||||
std::snprintf(buf, sizeof(buf), "%.0f%%", v * 100.f);
|
||||
return buf;
|
||||
}
|
||||
|
||||
static std::string fmt_sim(float v) {
|
||||
char buf[8];
|
||||
std::snprintf(buf, sizeof(buf), "%.2f", v);
|
||||
return buf;
|
||||
}
|
||||
|
||||
static std::string sanitise(const std::string& s) {
|
||||
std::string out;
|
||||
out.reserve(s.size());
|
||||
for (char c : s)
|
||||
out += (std::isalnum(c) ? std::tolower(c) : '_');
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // SAE_DEBUG
|
||||
@@ -0,0 +1,66 @@
|
||||
#pragma once
|
||||
#include "arcface_embedder.hpp"
|
||||
#include "config.hpp"
|
||||
#include "ort_provider.hpp"
|
||||
#include "trt_arcface_embedder.hpp"
|
||||
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
// ── EmbedderFunc ──────────────────────────────────────────────────────────────
|
||||
// KPN node: runs ArcFace on every 112×112 crop in an AlignedSceneFrame,
|
||||
// producing one L2-normalised 512-dim embedding per face.
|
||||
//
|
||||
// Backend selection:
|
||||
// --arcface-engine <path> → TrtArcFaceEmbedder (raw TensorRT, no ORT)
|
||||
// otherwise → ArcFaceEmbedder (ONNX Runtime, picks best EP)
|
||||
//
|
||||
// All crops in one frame are batched into a single forward pass (capped at
|
||||
// embed_batch_size). The backends serialise themselves; we only call them
|
||||
// from the single embedder thread.
|
||||
|
||||
struct EmbedderFunc {
|
||||
static constexpr std::string_view label() { return "embedder"; }
|
||||
|
||||
explicit EmbedderFunc(const Config& cfg, OrtProvider provider)
|
||||
: batch_size_(std::max(1, cfg.embed_batch_size))
|
||||
{
|
||||
if (!cfg.arcface_engine.empty()) {
|
||||
trt_ = std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
|
||||
if (trt_->max_batch() < static_cast<int>(batch_size_))
|
||||
throw std::runtime_error(
|
||||
"embed_batch_size " + std::to_string(batch_size_) +
|
||||
" exceeds engine max_batch " + std::to_string(trt_->max_batch()) +
|
||||
" — rebuild engine with EMBED_BATCH=" + std::to_string(batch_size_));
|
||||
} else {
|
||||
ort_ = std::make_unique<ArcFaceEmbedder>(
|
||||
cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
|
||||
}
|
||||
}
|
||||
|
||||
EmbeddedSceneFrame operator()(AlignedSceneFrame af) {
|
||||
if (af.source.eof || af.crops.empty())
|
||||
return {std::move(af.source), {}, {}, {}};
|
||||
|
||||
const auto& crops = af.crops;
|
||||
std::vector<Embedding> embeddings;
|
||||
embeddings.reserve(crops.size());
|
||||
for (size_t i = 0; i < crops.size(); i += batch_size_) {
|
||||
const size_t end = std::min(i + batch_size_, crops.size());
|
||||
std::vector<cv::Mat> chunk_crops(crops.begin() + i, crops.begin() + end);
|
||||
auto chunk = trt_ ? trt_->embed(chunk_crops) : ort_->embed(chunk_crops);
|
||||
embeddings.insert(embeddings.end(), chunk.begin(), chunk.end());
|
||||
}
|
||||
|
||||
return {std::move(af.source),
|
||||
std::move(af.faces),
|
||||
std::move(af.crops),
|
||||
std::move(embeddings)};
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<ArcFaceEmbedder> ort_;
|
||||
std::unique_ptr<TrtArcFaceEmbedder> trt_;
|
||||
size_t batch_size_;
|
||||
};
|
||||
@@ -0,0 +1,39 @@
|
||||
#pragma once
|
||||
#include "face_utils.hpp"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
// ── FaceAlignerFunc ───────────────────────────────────────────────────────────
|
||||
// KPN node: applies a 5-point similarity transform to each detected face,
|
||||
// producing a 112×112 BGR crop suitable for ArcFace inference.
|
||||
//
|
||||
// Alignment uses cv::estimateAffinePartial2D (RANSAC) to fit the detected
|
||||
// landmarks to ArcFace canonical positions. Degenerate detections (where the
|
||||
// affine fit fails) are silently dropped from the output vectors.
|
||||
|
||||
struct FaceAlignerFunc {
|
||||
static constexpr std::string_view label() { return "face_aligner"; }
|
||||
|
||||
AlignedSceneFrame operator()(SceneFrame sf) {
|
||||
if (sf.source.eof || sf.faces.empty())
|
||||
return {std::move(sf.source), {}, {}};
|
||||
|
||||
std::vector<DetectedFace> good_faces;
|
||||
std::vector<cv::Mat> crops;
|
||||
good_faces.reserve(sf.faces.size());
|
||||
crops.reserve(sf.faces.size());
|
||||
|
||||
for (auto& face : sf.faces) {
|
||||
cv::Mat crop = align_face(sf.source.image, face.landmarks);
|
||||
if (crop.empty()) {
|
||||
std::cerr << "[face_aligner] degenerate detection skipped\n";
|
||||
continue;
|
||||
}
|
||||
good_faces.push_back(face);
|
||||
crops.push_back(std::move(crop));
|
||||
}
|
||||
|
||||
return {std::move(sf.source), std::move(good_faces), std::move(crops)};
|
||||
}
|
||||
|
||||
};
|
||||
@@ -0,0 +1,61 @@
|
||||
#pragma once
|
||||
#include "scrfd_decoder.hpp"
|
||||
#include "trt_scrfd_decoder.hpp"
|
||||
#include "config.hpp"
|
||||
#include "ort_provider.hpp"
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
// ── FaceDetectorFunc ──────────────────────────────────────────────────────────
|
||||
// KPN node: runs SCRFD-500MF to detect ALL faces in a frame.
|
||||
//
|
||||
// Backend selection:
|
||||
// --detector-engine <path> → TrtScrfdDecoder (raw TensorRT, no ORT)
|
||||
// otherwise → SCRFDDecoder (ONNX Runtime)
|
||||
|
||||
struct FaceDetectorFunc {
|
||||
static constexpr std::string_view label() { return "face_detector"; }
|
||||
|
||||
explicit FaceDetectorFunc(const Config& cfg, OrtProvider provider)
|
||||
: max_faces_(cfg.max_faces)
|
||||
, min_face_px_(cfg.min_face_px)
|
||||
{
|
||||
if (!cfg.detector_engine.empty()) {
|
||||
trt_ = std::make_unique<TrtScrfdDecoder>(
|
||||
cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
|
||||
} else {
|
||||
ort_ = std::make_unique<SCRFDDecoder>(
|
||||
cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
|
||||
}
|
||||
}
|
||||
|
||||
SceneFrame operator()(Frame f) {
|
||||
if (f.eof) return {std::move(f), {}};
|
||||
|
||||
auto faces = trt_ ? trt_->detect(f.image) : ort_->detect(f.image);
|
||||
|
||||
// Drop faces below minimum pixel size (too small for reliable ArcFace alignment)
|
||||
faces.erase(
|
||||
std::remove_if(faces.begin(), faces.end(), [&](const DetectedFace& d) {
|
||||
return d.bbox.width < min_face_px_ || d.bbox.height < min_face_px_;
|
||||
}),
|
||||
faces.end());
|
||||
|
||||
// Sort largest-first so max_faces_ keeps the most informative detections
|
||||
std::sort(faces.begin(), faces.end(),
|
||||
[](const DetectedFace& a, const DetectedFace& b) {
|
||||
return a.bbox.area() > b.bbox.area();
|
||||
});
|
||||
if (static_cast<int>(faces.size()) > max_faces_)
|
||||
faces.resize(max_faces_);
|
||||
|
||||
return {std::move(f), std::move(faces)};
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<SCRFDDecoder> ort_;
|
||||
std::unique_ptr<TrtScrfdDecoder> trt_;
|
||||
int max_faces_{10};
|
||||
float min_face_px_{40.f};
|
||||
};
|
||||
@@ -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_;
|
||||
};
|
||||
@@ -0,0 +1,144 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
#include "ffmpeg_decoder.hpp"
|
||||
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <chrono>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
|
||||
// ── FrameSourceFunc ───────────────────────────────────────────────────────────
|
||||
// KPN source node: reads a movie file and emits one Frame per sample interval.
|
||||
//
|
||||
// Decode backend: FFmpeg with NVDEC (_cuvid) when available, CPU otherwise.
|
||||
//
|
||||
// Sampling strategy: seek to the next target timestamp rather than decoding
|
||||
// every frame, which is fast even for 1-FPS sampling of a 2-hour film.
|
||||
//
|
||||
// EOF handling: when the movie ends, emits a Frame with eof=true, then sleeps
|
||||
// 500 ms between subsequent calls until the KPN network stops the thread.
|
||||
|
||||
struct FrameSourceFunc {
|
||||
static constexpr std::string_view label() { return "frame_source"; }
|
||||
|
||||
explicit FrameSourceFunc(const Config& cfg)
|
||||
: decoder_(std::make_unique<FFmpegDecoder>(cfg.movie_path))
|
||||
{
|
||||
sample_interval_sec_ = 1.0 / cfg.sample_fps;
|
||||
next_pos_sec_ = cfg.start_sec;
|
||||
end_sec_ = cfg.end_sec;
|
||||
cut_threshold_ = cfg.cut_threshold;
|
||||
max_decode_fps_ = cfg.max_decode_fps;
|
||||
|
||||
double total_s = decoder_->duration_sec();
|
||||
double span_s = (end_sec_ > 0 ? std::min(end_sec_, total_s) : total_s)
|
||||
- cfg.start_sec;
|
||||
int n_frames = static_cast<int>(span_s * cfg.sample_fps);
|
||||
std::cerr << "[frame_source] decoder=" << decoder_->codec_name()
|
||||
<< " (" << (decoder_->hw_active() ? "NVDEC" : "CPU") << ")"
|
||||
<< " video_fps=" << decoder_->fps()
|
||||
<< " start=" << cfg.start_sec << "s"
|
||||
<< (end_sec_ > 0 ? " end=" + std::to_string(end_sec_) + "s" : "")
|
||||
<< " sample_fps=" << cfg.sample_fps
|
||||
<< " frames_to_emit=" << n_frames << "\n";
|
||||
}
|
||||
|
||||
Frame operator()() {
|
||||
if (hit_eof_) {
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(500));
|
||||
return Frame{{}, 0.0, -1, /*eof=*/true};
|
||||
}
|
||||
|
||||
// Wall-clock rate cap. KPN source nodes resubmit immediately on push
|
||||
// overflow, with no backpressure; without this cap we'd decode-and-drop
|
||||
// in a tight loop whenever downstream stalls. The cap also protects
|
||||
// ORT-only deployments where the pipeline can't keep up at decode speed.
|
||||
if (max_decode_fps_ > 0.f) {
|
||||
const auto now = std::chrono::steady_clock::now();
|
||||
if (!rate_started_) {
|
||||
rate_started_ = true;
|
||||
next_decode_at_ = now;
|
||||
}
|
||||
if (now < next_decode_at_)
|
||||
std::this_thread::sleep_until(next_decode_at_);
|
||||
const auto period = std::chrono::nanoseconds(
|
||||
static_cast<int64_t>(1e9f / max_decode_fps_));
|
||||
// Anchor the next slot off the slot we just consumed, not off
|
||||
// wall-clock now() — keeps the average rate stable. If we fell
|
||||
// behind by more than one period, snap forward to avoid building
|
||||
// up an unbounded sleep debt.
|
||||
next_decode_at_ += period;
|
||||
if (next_decode_at_ < now)
|
||||
next_decode_at_ = now + period;
|
||||
}
|
||||
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
cv::Mat img = decoder_->read_at(next_pos_sec_);
|
||||
auto t1 = std::chrono::steady_clock::now();
|
||||
double decode_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
|
||||
decode_ms_acc_ += decode_ms;
|
||||
++decode_count_;
|
||||
if (decode_count_ % 10 == 0) {
|
||||
double avg_ms = decode_ms_acc_ / 10.0;
|
||||
double avg_fps = avg_ms > 0.0 ? 1000.0 / avg_ms : 0.0;
|
||||
std::cerr << "[frame_source] decode avg=" << avg_ms << "ms"
|
||||
<< " fps=" << avg_fps << "\n";
|
||||
decode_ms_acc_ = 0.0;
|
||||
}
|
||||
|
||||
if (img.empty()) {
|
||||
hit_eof_ = true;
|
||||
std::cerr << "[frame_source] EOF at t=" << next_pos_sec_ << "s\n";
|
||||
return Frame{{}, next_pos_sec_, frame_idx_++, /*eof=*/true};
|
||||
}
|
||||
|
||||
// Cut detection: compare grayscale histogram to previous frame
|
||||
bool is_cut = false;
|
||||
cv::Mat gray;
|
||||
cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);
|
||||
cv::Mat hist;
|
||||
const int bins = 64;
|
||||
const float range[] = {0.f, 256.f};
|
||||
const float* ranges = range;
|
||||
cv::calcHist(&gray, 1, nullptr, cv::Mat(), hist, 1, &bins, &ranges);
|
||||
cv::normalize(hist, hist, 1.0, 0.0, cv::NORM_L1);
|
||||
|
||||
if (prev_hist_valid_) {
|
||||
double corr = cv::compareHist(prev_hist_, hist, cv::HISTCMP_CORREL);
|
||||
is_cut = (corr < cut_threshold_);
|
||||
if (is_cut)
|
||||
std::cerr << "[frame_source] cut at t=" << next_pos_sec_
|
||||
<< "s hist_corr=" << corr << "\n";
|
||||
}
|
||||
prev_hist_ = hist;
|
||||
prev_hist_valid_ = true;
|
||||
|
||||
Frame f{img, next_pos_sec_, frame_idx_++, /*eof=*/false, is_cut};
|
||||
next_pos_sec_ += sample_interval_sec_;
|
||||
if (end_sec_ > 0 && next_pos_sec_ > end_sec_) {
|
||||
hit_eof_ = true;
|
||||
std::cerr << "[frame_source] reached end_sec=" << end_sec_ << "s\n";
|
||||
}
|
||||
return f;
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<FFmpegDecoder> decoder_;
|
||||
double sample_interval_sec_{1.0};
|
||||
double next_pos_sec_{0.0};
|
||||
double end_sec_{-1.0};
|
||||
float cut_threshold_{0.70f};
|
||||
float max_decode_fps_{0.f};
|
||||
std::chrono::steady_clock::time_point next_decode_at_{};
|
||||
bool rate_started_{false};
|
||||
int64_t frame_idx_{0};
|
||||
bool hit_eof_{false};
|
||||
cv::Mat prev_hist_;
|
||||
bool prev_hist_valid_{false};
|
||||
double decode_ms_acc_{0.0};
|
||||
int decode_count_{0};
|
||||
};
|
||||
@@ -0,0 +1,154 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
#include "gallery/gallery_store.hpp"
|
||||
#include "gallery/gallery_calibration.hpp"
|
||||
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <iostream>
|
||||
|
||||
// ── 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 all
|
||||
// reference embeddings and taking the closest (best-of-N).
|
||||
|
||||
struct IdentityMatcherFunc {
|
||||
static constexpr std::string_view label() { return "identity_matcher"; }
|
||||
|
||||
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)
|
||||
{
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
cal_ = calibrate_gallery(flat_emb_, flat_actor_);
|
||||
|
||||
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";
|
||||
}
|
||||
|
||||
MatchedSceneFrame operator()(TrackedSceneFrame tf) {
|
||||
if (tf.source.eof) return {std::move(tf.source), {}};
|
||||
|
||||
std::vector<IdentifiedActor> actors;
|
||||
actors.reserve(tf.embeddings.size());
|
||||
|
||||
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];
|
||||
|
||||
// Per-actor best cosine similarity (max dot product)
|
||||
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]);
|
||||
int ai = flat_actor_[ei];
|
||||
if (sim > best_sim[ai]) best_sim[ai] = sim;
|
||||
}
|
||||
|
||||
// Find best and second-best actor by similarity
|
||||
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.similarity = cal_.valid
|
||||
? cal_.probability(best_s, log_prior_odds_)
|
||||
: best_s;
|
||||
}
|
||||
// actor_idx == -1, name == "" → unknown face
|
||||
|
||||
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_;
|
||||
};
|
||||
@@ -0,0 +1,140 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <kpn/main_thread_node.hpp>
|
||||
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
|
||||
#include <string>
|
||||
|
||||
// ── PreviewNode ───────────────────────────────────────────────────────────────
|
||||
// MainThreadNode: receives MatchedSceneFrame, draws annotations and shows the
|
||||
// frame in an OpenCV window. Runs on the main thread via preview.step().
|
||||
//
|
||||
// Returns false on EOF (ends the main loop) or when 'q' / Escape is pressed.
|
||||
|
||||
class PreviewNode : public kpn::MainThreadNode<PreviewNode,
|
||||
kpn::in<"matched">,
|
||||
MatchedSceneFrame> {
|
||||
public:
|
||||
static constexpr std::string_view label() { return "preview"; }
|
||||
explicit PreviewNode(const Config& cfg, std::size_t fifo_capacity = 4)
|
||||
: kpn::MainThreadNode<PreviewNode, kpn::in<"matched">, MatchedSceneFrame>(fifo_capacity)
|
||||
, max_display_w_(1280)
|
||||
{
|
||||
cv::namedWindow("scene_preview", cv::WINDOW_NORMAL);
|
||||
cv::resizeWindow("scene_preview", 1280, 720);
|
||||
}
|
||||
|
||||
// Called by step() on the main thread for each ready MatchedSceneFrame.
|
||||
bool operator()(MatchedSceneFrame mf) {
|
||||
if (mf.source.eof) return false;
|
||||
|
||||
cv::Mat display = mf.source.image.clone();
|
||||
draw_detections(display, mf.actors);
|
||||
draw_hud(display, mf.source.timestamp_sec, mf.actors);
|
||||
|
||||
// Fit to display width while keeping aspect ratio
|
||||
if (display.cols > max_display_w_) {
|
||||
float scale = static_cast<float>(max_display_w_) / display.cols;
|
||||
cv::resize(display, display, {}, scale, scale, cv::INTER_AREA);
|
||||
}
|
||||
|
||||
cv::imshow("scene_preview", display);
|
||||
|
||||
int key = cv::waitKey(1) & 0xFF;
|
||||
return (key != 'q' && key != 27 /* Esc */);
|
||||
}
|
||||
|
||||
private:
|
||||
int max_display_w_;
|
||||
|
||||
static void draw_detections(cv::Mat& img,
|
||||
const std::vector<IdentifiedActor>& actors) {
|
||||
for (const auto& ia : actors) {
|
||||
bool known = (ia.actor_idx >= 0);
|
||||
|
||||
// Green for identified, orange for unknown
|
||||
cv::Scalar box_colour = known
|
||||
? cv::Scalar(0, 210, 60)
|
||||
: cv::Scalar(0, 140, 255);
|
||||
|
||||
// Scale bbox to display image
|
||||
cv::Rect2f b = ia.bbox;
|
||||
cv::rectangle(img, b, box_colour, 2, cv::LINE_AA);
|
||||
|
||||
std::string tid = (ia.track_id >= 0) ? (" #" + std::to_string(ia.track_id)) : "";
|
||||
std::string label = known
|
||||
? (ia.name + tid + " " + pct(ia.similarity))
|
||||
: ("?" + tid);
|
||||
|
||||
// Dark backing strip so text is readable on any background
|
||||
int baseline = 0;
|
||||
cv::Size ts = cv::getTextSize(label, cv::FONT_HERSHEY_DUPLEX,
|
||||
0.55, 1, &baseline);
|
||||
cv::Point tl(static_cast<int>(b.x),
|
||||
std::max(0, static_cast<int>(b.y) - ts.height - 6));
|
||||
cv::Rect backing(tl.x, tl.y, ts.width + 8, ts.height + 8);
|
||||
backing &= cv::Rect(0, 0, img.cols, img.rows);
|
||||
if (backing.area() > 0)
|
||||
cv::rectangle(img, backing, box_colour * 0.6, cv::FILLED);
|
||||
|
||||
cv::putText(img, label,
|
||||
cv::Point(tl.x + 4, tl.y + ts.height + 2),
|
||||
cv::FONT_HERSHEY_DUPLEX, 0.55,
|
||||
cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
|
||||
}
|
||||
}
|
||||
|
||||
static void draw_hud(cv::Mat& img, double timestamp_sec,
|
||||
const std::vector<IdentifiedActor>& actors) {
|
||||
int known = 0;
|
||||
int unknown = 0;
|
||||
for (const auto& a : actors) (a.actor_idx >= 0 ? known : unknown)++;
|
||||
|
||||
// Timestamp and actor count in top-left corner
|
||||
char buf[128];
|
||||
std::snprintf(buf, sizeof(buf),
|
||||
"t = %dm %02ds | %d identified %d unknown",
|
||||
static_cast<int>(timestamp_sec) / 60,
|
||||
static_cast<int>(timestamp_sec) % 60,
|
||||
known, unknown);
|
||||
|
||||
int baseline = 0;
|
||||
cv::Size ts = cv::getTextSize(buf, cv::FONT_HERSHEY_SIMPLEX,
|
||||
0.6, 1, &baseline);
|
||||
cv::Rect hud(0, 0, ts.width + 16, ts.height + 12);
|
||||
hud &= cv::Rect(0, 0, img.cols, img.rows);
|
||||
cv::rectangle(img, hud, cv::Scalar(20, 20, 20), cv::FILLED);
|
||||
cv::putText(img, buf, cv::Point(8, ts.height + 6),
|
||||
cv::FONT_HERSHEY_SIMPLEX, 0.6,
|
||||
cv::Scalar(220, 220, 220), 1, cv::LINE_AA);
|
||||
|
||||
// Active actor name strip along the bottom
|
||||
if (known > 0) {
|
||||
std::string names;
|
||||
for (const auto& a : actors) {
|
||||
if (a.actor_idx < 0) continue;
|
||||
if (!names.empty()) names += " ";
|
||||
names += a.name;
|
||||
}
|
||||
cv::Size ns = cv::getTextSize(names, cv::FONT_HERSHEY_SIMPLEX,
|
||||
0.55, 1, &baseline);
|
||||
int y = img.rows - ns.height - 10;
|
||||
cv::Rect strip(0, y - 6, img.cols, ns.height + 16);
|
||||
strip &= cv::Rect(0, 0, img.cols, img.rows);
|
||||
cv::rectangle(img, strip, cv::Scalar(10, 10, 10, 180), cv::FILLED);
|
||||
cv::putText(img, names, cv::Point(8, img.rows - 10),
|
||||
cv::FONT_HERSHEY_SIMPLEX, 0.55,
|
||||
cv::Scalar(80, 220, 100), 1, cv::LINE_AA);
|
||||
}
|
||||
}
|
||||
|
||||
static std::string pct(float v) {
|
||||
char buf[8];
|
||||
std::snprintf(buf, sizeof(buf), "%.0f%%", v * 100.f);
|
||||
return buf;
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,204 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
using json = nlohmann::json;
|
||||
|
||||
// ── ResultSinkFunc ────────────────────────────────────────────────────────────
|
||||
// KPN sink node: accumulates SceneAnnotations and writes the final JSON on EOF.
|
||||
//
|
||||
// Verbosity::minimal — merges per-frame presence into contiguous time windows.
|
||||
// Output: { "movie": "...", "actors": [{ "name", "imdb_id", "scenes": [[t0,t1], ...] }] }
|
||||
//
|
||||
// Verbosity::standard — per-frame detail including bboxes, similarity, unknowns.
|
||||
// Output: { "frames": [{ "t", "identified": [...], "unknowns": [...] }] }
|
||||
//
|
||||
// eof signal: sets done_ = true so the main thread can call net.stop().
|
||||
|
||||
struct ResultSinkFunc {
|
||||
static constexpr std::string_view label() { return "result_sink"; }
|
||||
|
||||
ResultSinkFunc(const Config& cfg, std::atomic<bool>& done)
|
||||
: cfg_(cfg), done_(done)
|
||||
{}
|
||||
|
||||
void operator()(SceneAnnotation sa) {
|
||||
if (sa.eof) {
|
||||
flush();
|
||||
return;
|
||||
}
|
||||
|
||||
// Progress to stderr
|
||||
std::cerr << "\r[result_sink] t=" << sa.timestamp_sec << "s"
|
||||
<< " active=" << count_known(sa.visible_actors)
|
||||
<< " unknowns=" << count_unknown(sa.visible_actors)
|
||||
<< std::flush;
|
||||
|
||||
frames_.push_back(std::move(sa));
|
||||
}
|
||||
|
||||
// Write accumulated results and signal done. Safe to call more than once.
|
||||
void flush() {
|
||||
if (written_.exchange(true)) return;
|
||||
write_output();
|
||||
done_.store(true, std::memory_order_release);
|
||||
}
|
||||
|
||||
private:
|
||||
static int count_known(const std::vector<IdentifiedActor>& v) {
|
||||
int n = 0;
|
||||
for (const auto& a : v) if (a.actor_idx >= 0) ++n;
|
||||
return n;
|
||||
}
|
||||
static int count_unknown(const std::vector<IdentifiedActor>& v) {
|
||||
int n = 0;
|
||||
for (const auto& a : v) if (a.actor_idx < 0) ++n;
|
||||
return n;
|
||||
}
|
||||
|
||||
void write_output() {
|
||||
std::cerr << "\n[result_sink] writing " << cfg_.output_path << "\n";
|
||||
|
||||
json root;
|
||||
if (cfg_.verbosity == Verbosity::xray) {
|
||||
root = build_xray();
|
||||
} else {
|
||||
root["movie"] = cfg_.movie_path;
|
||||
root["sample_fps"] = cfg_.sample_fps;
|
||||
root["anneal_sec"] = cfg_.anneal_sec;
|
||||
root["actors"] = build_epochs();
|
||||
if (cfg_.verbosity == Verbosity::standard)
|
||||
root["frames"] = build_standard();
|
||||
}
|
||||
|
||||
std::ofstream f(cfg_.output_path);
|
||||
if (!f.is_open()) {
|
||||
std::cerr << "[result_sink] ERROR: cannot write " << cfg_.output_path << "\n";
|
||||
return;
|
||||
}
|
||||
f << root.dump(2) << "\n";
|
||||
std::cerr << "[result_sink] done.\n";
|
||||
}
|
||||
|
||||
struct ActorWindow {
|
||||
std::string name, imdb_id;
|
||||
std::vector<std::pair<double, double>> scenes; // [start_sec, end_sec]
|
||||
};
|
||||
|
||||
// Core logic: merge per-frame detections into annealed [start, end] windows.
|
||||
std::vector<ActorWindow> build_actor_windows() {
|
||||
struct Info { std::string name, imdb_id; };
|
||||
std::map<int, Info> actor_info;
|
||||
std::map<int, std::vector<double>> timestamps;
|
||||
|
||||
for (const auto& frame : frames_) {
|
||||
for (const auto& ia : frame.visible_actors) {
|
||||
if (ia.actor_idx < 0) continue;
|
||||
actor_info[ia.actor_idx] = {ia.name, ia.imdb_id};
|
||||
timestamps[ia.actor_idx].push_back(frame.timestamp_sec);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<ActorWindow> result;
|
||||
for (auto& [idx, ts_vec] : timestamps) {
|
||||
ActorWindow aw;
|
||||
aw.name = actor_info[idx].name;
|
||||
aw.imdb_id = actor_info[idx].imdb_id;
|
||||
|
||||
double win_start = ts_vec[0], win_end = ts_vec[0];
|
||||
for (size_t i = 1; i < ts_vec.size(); ++i) {
|
||||
if (ts_vec[i] - win_end > cfg_.anneal_sec) {
|
||||
aw.scenes.push_back({win_start, win_end});
|
||||
win_start = ts_vec[i];
|
||||
}
|
||||
win_end = ts_vec[i];
|
||||
}
|
||||
aw.scenes.push_back({win_start, win_end});
|
||||
result.push_back(std::move(aw));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
json build_epochs() {
|
||||
json actors = json::array();
|
||||
for (const auto& aw : build_actor_windows()) {
|
||||
json windows = json::array();
|
||||
for (const auto& [s, e] : aw.scenes)
|
||||
windows.push_back({s, e});
|
||||
json ja;
|
||||
ja["name"] = aw.name;
|
||||
ja["imdb_id"] = aw.imdb_id;
|
||||
ja["scenes"] = std::move(windows);
|
||||
actors.push_back(std::move(ja));
|
||||
}
|
||||
return actors;
|
||||
}
|
||||
|
||||
// Jellyfin-Xray format: { "second": ["Actor", ...] }
|
||||
// Expands each annealed window into every integer second so coverage is dense
|
||||
// regardless of sample rate. Seconds between scenes have no key → overlay clears.
|
||||
json build_xray() {
|
||||
std::map<int, std::vector<std::string>> xray;
|
||||
for (const auto& aw : build_actor_windows()) {
|
||||
for (const auto& [start, end] : aw.scenes) {
|
||||
int t0 = static_cast<int>(std::floor(start));
|
||||
int t1 = static_cast<int>(std::ceil(end));
|
||||
for (int t = t0; t <= t1; ++t)
|
||||
xray[t].push_back(aw.name);
|
||||
}
|
||||
}
|
||||
|
||||
json root = json::object();
|
||||
for (const auto& [t, names] : xray)
|
||||
root[std::to_string(t)] = names;
|
||||
return root;
|
||||
}
|
||||
|
||||
json build_standard() {
|
||||
json frames = json::array();
|
||||
for (const auto& frame : frames_) {
|
||||
json jf;
|
||||
jf["t"] = frame.timestamp_sec;
|
||||
jf["identified"] = json::array();
|
||||
jf["unknowns"] = json::array();
|
||||
|
||||
for (const auto& ia : frame.visible_actors) {
|
||||
const auto& b = ia.bbox;
|
||||
json jbox = {b.x, b.y, b.width, b.height};
|
||||
|
||||
if (ia.actor_idx >= 0) {
|
||||
json ja;
|
||||
ja["name"] = ia.name;
|
||||
ja["imdb_id"] = ia.imdb_id;
|
||||
ja["similarity"] = ia.similarity;
|
||||
ja["track_id"] = ia.track_id;
|
||||
ja["bbox"] = jbox;
|
||||
jf["identified"].push_back(std::move(ja));
|
||||
} else {
|
||||
json ju;
|
||||
ju["bbox"] = jbox;
|
||||
ju["track_id"] = ia.track_id;
|
||||
ju["confidence"] = ia.similarity; // reuse field; 0 for unknowns
|
||||
jf["unknowns"].push_back(std::move(ju));
|
||||
}
|
||||
}
|
||||
frames.push_back(std::move(jf));
|
||||
}
|
||||
return frames;
|
||||
}
|
||||
|
||||
const Config& cfg_;
|
||||
std::atomic<bool>& done_;
|
||||
std::atomic<bool> written_{false};
|
||||
std::vector<SceneAnnotation> frames_;
|
||||
};
|
||||
@@ -0,0 +1,92 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <map>
|
||||
#include <iostream>
|
||||
|
||||
// ── SceneTrackerFunc ──────────────────────────────────────────────────────────
|
||||
// KPN node: maintains an extinction-timer state machine per identified actor.
|
||||
//
|
||||
// On each MatchedSceneFrame:
|
||||
// 1. Update last_seen for every matched known actor.
|
||||
// 2. Expire actors whose last_seen is older than extinction_sec.
|
||||
// 3. Emit SceneAnnotation with all currently active (non-expired) actors,
|
||||
// including their most recently seen bbox and best similarity score.
|
||||
//
|
||||
// Unknown faces (actor_idx == -1) are passed through per-frame but are NOT
|
||||
// tracked across frames — each frame reports its own unknowns independently.
|
||||
|
||||
struct SceneTrackerFunc {
|
||||
static constexpr std::string_view label() { return "scene_tracker"; }
|
||||
|
||||
explicit SceneTrackerFunc(const Config& cfg)
|
||||
: extinction_sec_(cfg.extinction_sec)
|
||||
{
|
||||
std::cerr << "[scene_tracker] extinction_sec=" << extinction_sec_ << "\n";
|
||||
}
|
||||
|
||||
SceneAnnotation operator()(MatchedSceneFrame mf) {
|
||||
if (mf.source.eof) return {0.0, {}, /*eof=*/true};
|
||||
|
||||
double now = mf.source.timestamp_sec;
|
||||
|
||||
// Update known actors
|
||||
for (const auto& ia : mf.actors) {
|
||||
if (ia.actor_idx < 0) continue; // skip unknowns
|
||||
|
||||
auto& slot = active_[ia.actor_idx];
|
||||
slot.last_seen = now;
|
||||
slot.last_bbox = ia.bbox;
|
||||
slot.last_crop = ia.crop;
|
||||
slot.name = ia.name;
|
||||
slot.imdb_id = ia.imdb_id;
|
||||
// Keep the best (highest) similarity seen in this window
|
||||
if (ia.similarity > slot.best_similarity)
|
||||
slot.best_similarity = ia.similarity;
|
||||
}
|
||||
|
||||
// Expire stale actors
|
||||
for (auto it = active_.begin(); it != active_.end(); ) {
|
||||
if ((now - it->second.last_seen) > extinction_sec_)
|
||||
it = active_.erase(it);
|
||||
else
|
||||
++it;
|
||||
}
|
||||
|
||||
// Build annotation: active known actors
|
||||
std::vector<IdentifiedActor> visible;
|
||||
visible.reserve(active_.size() + mf.actors.size());
|
||||
|
||||
for (const auto& [actor_idx, slot] : active_) {
|
||||
IdentifiedActor ia;
|
||||
ia.actor_idx = actor_idx;
|
||||
ia.name = slot.name;
|
||||
ia.imdb_id = slot.imdb_id;
|
||||
ia.similarity = slot.best_similarity;
|
||||
ia.bbox = slot.last_bbox;
|
||||
ia.crop = slot.last_crop;
|
||||
visible.push_back(ia);
|
||||
}
|
||||
|
||||
// Append per-frame unknowns (actor_idx == -1) directly
|
||||
for (const auto& ia : mf.actors) {
|
||||
if (ia.actor_idx < 0) visible.push_back(ia);
|
||||
}
|
||||
|
||||
return {now, std::move(visible)};
|
||||
}
|
||||
|
||||
private:
|
||||
struct Slot {
|
||||
double last_seen{0.0};
|
||||
float best_similarity{0.f};
|
||||
cv::Rect2f last_bbox;
|
||||
cv::Mat last_crop;
|
||||
std::string name;
|
||||
std::string imdb_id;
|
||||
};
|
||||
|
||||
double extinction_sec_;
|
||||
std::map<int, Slot> active_; // actor_idx → state
|
||||
};
|
||||
Reference in New Issue
Block a user