From 26139ffe8aef1f19980345f28629e56a125da4cf Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sun, 19 Jul 2026 19:05:05 +0200 Subject: [PATCH] feat(engine): add Python replay bindings, gallery pose-expansion, scene detection, embedding dumps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New C++ sources: - kpn_bindings.cpp (sae_kpn): assembles the real face_tracker/identity_matcher/ scene_tracker nodes inside a Python-driven KPN network via nanobind, for offline threshold-sweep replay against dumped embeddings (scripts/optimizer/). - track_gallery.hpp: per-film gallery expansion — promotes a confidently- identified track's novel-pose reference views into an in-memory annex so later frames/tracks of that actor at similar poses are recognised, without touching the baked gallery. - dump_embeddings.cpp: standalone exe that runs detect→embed only (no gallery, no matching) and dumps per-frame face embeddings + metadata to HDF5, so a parameter sweep can replay the expensive half once and vary tracking/matching config freely downstream. - scene_detector.hpp / scene_detector_node.hpp: TransNetV2-based shot-boundary detection, opt-in alongside the always-on histogram cut detector. - camera_position_change_detector_node.hpp, embedding_dump_node.hpp: supporting nodes for the above. --- src/dump_embeddings.cpp | 92 +++++++ src/gallery/track_gallery.hpp | 243 +++++++++++++++++ src/inference/scene_detector.hpp | 45 +++ src/kpn_bindings.cpp | 257 ++++++++++++++++++ .../camera_position_change_detector_node.hpp | 73 +++++ src/nodes/embedding_dump_node.hpp | 124 +++++++++ src/nodes/scene_detector_node.hpp | 179 ++++++++++++ 7 files changed, 1013 insertions(+) create mode 100644 src/dump_embeddings.cpp create mode 100644 src/gallery/track_gallery.hpp create mode 100644 src/inference/scene_detector.hpp create mode 100644 src/kpn_bindings.cpp create mode 100644 src/nodes/camera_position_change_detector_node.hpp create mode 100644 src/nodes/embedding_dump_node.hpp create mode 100644 src/nodes/scene_detector_node.hpp diff --git a/src/dump_embeddings.cpp b/src/dump_embeddings.cpp new file mode 100644 index 0000000..bce34b7 --- /dev/null +++ b/src/dump_embeddings.cpp @@ -0,0 +1,92 @@ +// dump_embeddings — standalone embedding dumper (NO gallery required). +// +// Runs only the expensive, gallery-independent front half of the pipeline +// (decode → camera-pos → detect → align → embed) and writes per-frame face +// embeddings + metadata to HDF5 (scripts/optimizer/SCHEMA.md). Unlike +// `scene_analyze --dump-embeddings`, it does NOT construct the identity matcher, so +// it neither loads a gallery nor runs calibration — ~24s faster per run and no +// gallery file needed. Purpose-built for the optimizer's replay corpus and the +// embedding-model bake-off (dump each --arcface model over the film set). +// +// Usage: +// dump_embeddings --movie --out [--arcface ] +// [--detector ] [--fps 1] [--start S] [--end S] +// [--conf 0.5] [--max-faces 10] [--min-face-px 40] + +#include "config.hpp" +#include "types.hpp" +#include "nodes/frame_source_node.hpp" +#include "nodes/camera_position_change_detector_node.hpp" +#include "nodes/face_detector_node.hpp" +#include "nodes/face_aligner_node.hpp" +#include "nodes/embedder_node.hpp" +#include "nodes/embedding_dump_node.hpp" + +#include + +#include +#include +#include + +int main(int argc, char** argv) { + Config cfg; + cfg.arcface_model = kDefaultArcfaceModel; + cfg.detector_model = std::string(SAE_MODELS_DIR) + "/scrfd_500m_bnkps.onnx"; + + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + auto next = [&]() -> std::string { + if (++i >= argc) throw std::runtime_error("missing arg after " + a); + return argv[i]; + }; + if (a == "--movie") cfg.movie_path = next(); + else if (a == "--out") cfg.dump_embeddings_path = next(); + else if (a == "--arcface") cfg.arcface_model = next(); + else if (a == "--arcface-engine") cfg.arcface_engine = next(); + else if (a == "--detector") cfg.detector_model = next(); + else if (a == "--detector-engine") cfg.detector_engine = next(); + else if (a == "--fps") cfg.sample_fps = std::stof(next()); + else if (a == "--start") cfg.start_sec = std::stod(next()); + else if (a == "--end") cfg.end_sec = std::stod(next()); + else if (a == "--conf") cfg.detector_conf = std::stof(next()); + else if (a == "--max-faces") cfg.max_faces = std::stoi(next()); + else if (a == "--min-face-px") cfg.min_face_px = std::stof(next()); + else if (a == "--max-decode-fps") cfg.max_decode_fps = std::stof(next()); + else { std::cerr << "[dump] unknown flag: " << a << "\n"; return 1; } + } + if (cfg.movie_path.empty() || cfg.dump_embeddings_path.empty()) { + std::cerr << "Usage: dump_embeddings --movie --out " + "[--arcface ] [--fps 1] ...\n"; + return 1; + } + + std::atomic done{false}; + + FrameSourceFunc source_fn {cfg}; + CameraPositionChangeDetectorFunc campos_fn {cfg}; + FaceDetectorFunc detector_fn{cfg}; + FaceAlignerFunc aligner_fn; + EmbedderFunc embedder_fn{cfg}; + EmbeddingDumpFunc dump_fn {cfg, done}; + + kpn::ObjectNode, kpn::out<"raw">, "frame_source", 0> source (source_fn, 32); + kpn::ObjectNode, kpn::out<"frame">, "camera_pos", 0> campos (campos_fn, 32); + kpn::ObjectNode, kpn::out<"scene">, "face_detector", 0> detector(detector_fn, 64); + kpn::ObjectNode, kpn::out<"aligned">, "face_aligner", 0> aligner (aligner_fn, 64); + kpn::ObjectNode, kpn::out<"embedded">, "embedder", 0> embedder(embedder_fn, 32); + kpn::ObjectNode,kpn::out<>, "embedding_dump",0> dump (dump_fn, 32); + + auto net = kpn::make_network( + kpn::edge(source.output<"raw">(), campos.input<"raw">()), + kpn::edge(campos.output<"frame">(), detector.input<"frame">()), + kpn::edge(detector.output<"scene">(), aligner.input<"scene">()), + kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()), + kpn::edge(embedder.output<"embedded">(), dump.input<"embedded">()) + ); + + net.start(); + using namespace std::chrono_literals; + while (!done.load(std::memory_order_acquire)) std::this_thread::sleep_for(50ms); + net.stop(); + return 0; +} diff --git a/src/gallery/track_gallery.hpp b/src/gallery/track_gallery.hpp new file mode 100644 index 0000000..674f628 --- /dev/null +++ b/src/gallery/track_gallery.hpp @@ -0,0 +1,243 @@ +#pragma once +#include "types.hpp" +#include "config.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#ifdef SAE_DEBUG +#include +#endif +#include +#include + +// ── TrackGallery ────────────────────────────────────────────────────────────── +// Per-film gallery expansion driven by track continuity. +// +// A single uncut face track is, by construction, one physical person: the +// face_tracker links detections frame-to-frame and clears all tracks on a scene +// cut, so a track ID never spans a cut. That continuity is a same-identity label +// the baked gallery does not have. TrackGallery exploits it in two stages: +// +// 1. Diversity buffer (per track). Every frame's embedding is offered to a +// fixed-capacity buffer for that track. When full, the member whose best +// similarity to the *owning actor's* gallery references is HIGHEST is +// dropped — i.e. the pose the gallery already recognises well is the least +// informative, so the buffer is continuously biased toward the gallery-far +// (novel-pose) embeddings on which recognition currently fails. +// +// 2. Promotion (on confirmation). A track is "owned" by actor A once ≥N frames +// have been accepted (by the matcher's calibrated posterior) as A. On +// confirmation the retained buffer — the hard, gallery-far poses — is +// promoted into A's per-film annex, after two safety gates: +// • novelty: only embeddings whose best sim to A's refs is below +// expand_novelty_sim are added (skip poses already covered); +// • spread: if the retained buffer's internal spread (1 − min pairwise +// cosine sim) exceeds expand_track_spread_max the whole track is +// rejected — such spread signals a track-ID collision merging two +// people, whose embeddings must never enter A's annex. +// +// The annex is CPU-side and in-memory: it is small (tens of embeddings) so the +// matcher scans it with a scalar loop, and it is discarded when the process +// exits. Promoted embeddings only help SUBSEQUENT frames and later tracks of A — +// the pipeline stays streaming, no emitted output is buffered or relabelled. + +struct TrackGallery { + // One promoted reference view held in the per-actor annex. + struct AnnexEntry { + Embedding emb; + int actor_idx{-1}; + }; + + explicit TrackGallery(const Config& cfg) + : enabled_(cfg.expand_gallery) + , buffer_size_(std::max(1, cfg.expand_buffer_size)) + , novelty_sim_(cfg.expand_novelty_sim) + , spread_max_(cfg.expand_track_spread_max) + , min_anchor_frames_(std::max(1, cfg.expand_min_anchor_frames)) + , debug_dir_(cfg.expand_debug_dir) + { + if (!enabled_) return; + std::cerr << "[track_gallery] per-film expansion ON" + << " buffer=" << buffer_size_ + << " novelty_sim<" << novelty_sim_ + << " spread_max=" << spread_max_ + << " min_anchor_frames=" << min_anchor_frames_; + if (!debug_dir_.empty()) { + std::filesystem::create_directories(debug_dir_); + std::cerr << " debug_dir=" << debug_dir_; + } + std::cerr << "\n"; + } + + bool enabled() const { return enabled_; } + + // Current annex contents (empty when disabled). The matcher scans these + // alongside the baked gallery so a promoted view can win best-of-N for its + // actor. Returned by const-ref; only grows, never reordered. + const std::vector& annex() const { return annex_; } + + // Offer one observed face to its track's diversity buffer. + // track_id : face_tracker track (−1 = untracked, ignored) + // emb : this frame's raw embedding + // best_actor : actor with the highest gallery similarity for this face + // best_gal_sim : that similarity (best sim to best_actor's baked+annex refs) + // accepted : true if the matcher accepted this face as best_actor + // crop : aligned crop, retained only when debug dumping is on + void observe(int track_id, const Embedding& emb, + int best_actor, float best_gal_sim, bool accepted, + const cv::Mat& crop) + { + if (!enabled_ || track_id < 0) return; + + TrackState& ts = tracks_[track_id]; + + // Vote toward ownership: only accepted frames name an actor, and a track + // that flip-flops between actors is ambiguous, so we tally per actor and + // pick the plurality winner at confirmation time. + if (accepted && best_actor >= 0) { + ts.actor_votes[best_actor]++; + ts.accepted_frames++; + } + + insert_into_buffer(ts, emb, best_gal_sim, crop); + + // Confirm and promote as soon as the anchor threshold is met, once. + if (!ts.promoted && ts.accepted_frames >= min_anchor_frames_) + promote(track_id, ts); + } + + // Drop a track's buffer when the face_tracker expires it or on a scene cut, + // so stale/cross-cut embeddings can never be promoted later. Called by the + // matcher when it observes a cut or track disappearance. + void forget(int track_id) { tracks_.erase(track_id); } + + // Drop every track buffer (scene cut / EOF). Mirrors face_tracker's clear. + void clear_tracks() { tracks_.clear(); } + +private: + struct BufEntry { + Embedding emb; + float gal_sim{0.f}; // best sim to owning actor's refs when observed + cv::Mat crop; // populated only when debug_dir_ set + }; + + struct TrackState { + std::vector buf; + std::map actor_votes; // actor_idx → accepted-frame count + int accepted_frames{0}; + bool promoted{false}; + }; + + void insert_into_buffer(TrackState& ts, const Embedding& emb, + float gal_sim, const cv::Mat& crop) + { + BufEntry e; + e.emb = emb; + e.gal_sim = gal_sim; + if (!debug_dir_.empty() && !crop.empty()) e.crop = crop.clone(); + + if (static_cast(ts.buf.size()) < buffer_size_) { + ts.buf.push_back(std::move(e)); + return; + } + + // Buffer full: evict the member the gallery recognises best (highest + // gal_sim) — least informative — but only if the newcomer is at least as + // novel. Keeping the most gallery-far views is the whole point. + int worst_i = -1; + float worst_sim = e.gal_sim; // newcomer's sim is the bar to beat + for (int i = 0; i < static_cast(ts.buf.size()); ++i) { + if (ts.buf[i].gal_sim > worst_sim) { + worst_sim = ts.buf[i].gal_sim; + worst_i = i; + } + } + // worst_i == −1 → every buffered view is already more novel than the + // newcomer; drop the newcomer instead of a better sample. + if (worst_i >= 0) ts.buf[worst_i] = std::move(e); + } + + void promote(int track_id, TrackState& ts) { + ts.promoted = true; // idempotent: never promote a track twice + + int actor = plurality_actor(ts); + if (actor < 0) return; + + // ── Safety gate: internal spread ───────────────────────────────────── + // A legitimate single-person track varies in pose but stays reasonably + // self-similar. Large spread signals two people merged under one track + // ID — reject the whole track rather than poison the actor's annex. + float spread = buffer_spread(ts.buf); + if (spread > spread_max_) { + std::cerr << "[track_gallery] track " << track_id + << " → actor " << actor + << " REJECTED (spread " << spread + << " > " << spread_max_ << ", likely ID collision)\n"; + return; + } + + int added = 0; + for (const auto& be : ts.buf) { + // ── Safety gate: novelty ───────────────────────────────────────── + // Skip poses the gallery already covers; only gallery-far views are + // worth the annex slot (and the extra per-frame scan cost). + if (be.gal_sim >= novelty_sim_) continue; + annex_.push_back({be.emb, actor}); + if (!debug_dir_.empty() && !be.crop.empty()) + dump_mugshot(track_id, actor, added, be); + ++added; + } + + std::cerr << "[track_gallery] track " << track_id + << " confirmed actor " << actor + << " (" << ts.accepted_frames << " accepted frames, spread " + << spread << ") — promoted " << added << "/" + << ts.buf.size() << " views; annex now " + << annex_.size() << "\n"; + } + + static int plurality_actor(const TrackState& ts) { + int best = -1, best_votes = 0; + for (const auto& [ai, v] : ts.actor_votes) { + if (v > best_votes) { best_votes = v; best = ai; } + } + return best; + } + + // Spread = 1 − min pairwise cosine similarity over the buffer (0 when <2). + static float buffer_spread(const std::vector& buf) { + float min_sim = std::numeric_limits::max(); + for (size_t i = 0; i < buf.size(); ++i) + for (size_t j = i + 1; j < buf.size(); ++j) + min_sim = std::min(min_sim, cosine_similarity(buf[i].emb, buf[j].emb)); + if (min_sim == std::numeric_limits::max()) return 0.f; + return 1.f - min_sim; + } + + void dump_mugshot(int track_id, int actor, int idx, const BufEntry& be) { +#ifdef SAE_DEBUG + char name[64]; + std::snprintf(name, sizeof(name), "trk%d_actor%d_%d_sim%.3f.jpg", + track_id, actor, idx, be.gal_sim); + cv::imwrite((std::filesystem::path(debug_dir_) / name).string(), be.crop); +#else + (void)track_id; (void)actor; (void)idx; (void)be; +#endif + } + + bool enabled_; + int buffer_size_; + float novelty_sim_; + float spread_max_; + int min_anchor_frames_; + std::string debug_dir_; + + std::map tracks_; + std::vector annex_; +}; diff --git a/src/inference/scene_detector.hpp b/src/inference/scene_detector.hpp new file mode 100644 index 0000000..d443fbf --- /dev/null +++ b/src/inference/scene_detector.hpp @@ -0,0 +1,45 @@ +#pragma once +#include +#include + +#include + +// ── ISceneDetector ──────────────────────────────────────────────────────────── +// Backend-agnostic shot-boundary (scene-cut) detector interface, mirroring +// IFaceDetector. The concrete implementation wraps TransNetV2 and is selected at +// compile time by CMake (SAE_INFERENCE_BACKEND): exactly one of +// backends/ort_backend.cpp or backends/trt_backend.cpp provides +// make_scene_detector(). +// +// TransNetV2 consumes a window of exactly kWindow consecutive frames, each +// downscaled to kFrameW×kFrameH RGB, and predicts a per-frame boundary +// probability. It is a DENSE model: the frames it sees must be consecutive at +// (near) native frame rate — 1-FPS sampled frames give meaningless output. +// Feeding is handled by the pipeline (dense decode + decimator); this interface +// only exposes the fixed-size window inference. + +struct Config; + +struct ISceneDetector { + virtual ~ISceneDetector() = default; + + // TransNetV2 fixed input contract (from the elya5/transnetv2 ONNX export): + // input "input" : float32 [1, 100, 27, 48, 3] (RGB, channels-last, 0-255) + // output "534" : float32 [1, 100, 1] single-frame boundary logits + // output "535" : float32 [1, 100, 1] "many-hot" auxiliary head (unused) + static constexpr int kWindow = 100; + static constexpr int kFrameH = 27; + static constexpr int kFrameW = 48; + + // Run one window of exactly kWindow frames. Each frame must already be + // kFrameW×kFrameH, BGR, CV_8UC3 (the backend handles BGR→RGB). Returns + // kWindow boundary probabilities in [0,1] (sigmoid of the primary head), + // one per input frame, in order. + virtual std::vector detect_window( + const std::vector& window) = 0; +}; + +// Construct the scene detector for the compiled-in backend. Reads +// cfg.scene_model / cfg.scene_engine and cfg.trt. Only called when scene +// detection is enabled (cfg.scene_detect). +std::unique_ptr make_scene_detector(const Config& cfg); diff --git a/src/kpn_bindings.cpp b/src/kpn_bindings.cpp new file mode 100644 index 0000000..c105b6b --- /dev/null +++ b/src/kpn_bindings.cpp @@ -0,0 +1,257 @@ +// sae_kpn — run the real downstream pipeline nodes (face_tracker, identity_matcher, +// scene_tracker) inside a Python-assembled KPN network, fed by a Python HDF5 replay +// source. Lets a parameter sweep re-run the exact C++ matching/tracking logic over +// dumped embeddings — no video decode, no GPU — with different Config knobs each run. +// +// Boundary types (cross the Python seam): +// EmbeddedSceneFrame IN (built by the Python replay source from HDF5 arrays) +// SceneAnnotation OUT (read by the Python sink → presence JSON) +// Intermediate types (TrackedSceneFrame, MatchedSceneFrame) flow C++→C++ only, but +// still need channel factories + converters registered so PyNetwork can wire them. + +#define KPN_BUILD_PYTHON +#include +#include + +#include "types.hpp" +#include "config.hpp" +#include "gallery/gallery_store.hpp" +#include "nodes/face_tracker_node.hpp" +#include "nodes/identity_matcher_node.hpp" +#include "nodes/scene_tracker_node.hpp" + +#include +#include +#include +#include +#include + +#include +#include + +namespace nb = nanobind; +using namespace nb::literals; + +// The variant spanning every type that flows on a channel in the replay chain. +using SaeVariant = std::variant; + +// ── Converters ───────────────────────────────────────────────────────────────── +// Only EmbeddedSceneFrame (in) and SceneAnnotation (out) actually cross the seam; +// the two intermediates get identity-ish stubs (never converted in practice) so the +// variant's converter map is complete. + +namespace kpn { + +// EmbeddedSceneFrame: built FROM Python (a dict of numpy arrays). to_python is a +// stub (the replay source only produces it; nothing reads it back). +template<> struct PythonConverter { + static constexpr const char* type_name = "EmbeddedSceneFrame"; + + static nb::object to_python(const EmbeddedSceneFrame&) { + // Not needed downstream; return None. (Kept total for map completeness.) + return nb::none(); + } + + static EmbeddedSceneFrame from_python(nb::object o) { + nb::dict d = nb::cast(o); + EmbeddedSceneFrame ef; + ef.source.timestamp_sec = nb::cast(d["timestamp_sec"]); + ef.source.frame_idx = d.contains("frame_idx") ? nb::cast(d["frame_idx"]) : -1; + ef.source.eof = d.contains("eof") ? nb::cast(d["eof"]) : false; + ef.source.is_cut = d.contains("is_cut") ? nb::cast(d["is_cut"]) : false; + if (ef.source.eof) return ef; + + // faces: (N,4) bbox, (N,10) landmarks, (N,) confidence, (N,512) embeddings + auto bbox = nb::cast, nb::c_contig>>(d["bbox"]); + auto lmk = nb::cast, nb::c_contig>>(d["landmarks"]); + auto conf = nb::cast, nb::c_contig>>(d["confidence"]); + auto emb = nb::cast, nb::c_contig>>(d["embeddings"]); + + const size_t n = bbox.shape(0); + ef.faces.reserve(n); + ef.embeddings.reserve(n); + const float* bp = bbox.data(); + const float* lp = lmk.data(); + const float* cp = conf.data(); + const float* ep = emb.data(); + for (size_t i = 0; i < n; ++i) { + DetectedFace f; + f.bbox = cv::Rect2f(bp[i*4+0], bp[i*4+1], bp[i*4+2], bp[i*4+3]); + for (int k = 0; k < 5; ++k) + f.landmarks[k] = cv::Point2f(lp[i*10 + k*2], lp[i*10 + k*2 + 1]); + f.confidence = cp[i]; + ef.faces.push_back(f); + + Embedding e; + for (int k = 0; k < 512; ++k) e[k] = ep[i*512 + k]; + ef.embeddings.push_back(e); + } + // crops left empty: tracker/matcher only forward them for debug rendering. + ef.crops.resize(n); + return ef; + } +}; + +// SceneAnnotation: read INTO Python. from_python is a stub (Python never builds one). +template<> struct PythonConverter { + static constexpr const char* type_name = "SceneAnnotation"; + + static nb::object to_python(const SceneAnnotation& sa) { + nb::dict d; + d["timestamp_sec"] = sa.timestamp_sec; + d["eof"] = sa.eof; + nb::list actors; + for (const auto& a : sa.visible_actors) { + nb::dict ad; + ad["actor_idx"] = a.actor_idx; + ad["track_id"] = a.track_id; + ad["name"] = a.name; + ad["imdb_id"] = a.imdb_id; + ad["tmdb_id"] = a.tmdb_id; + ad["jellyfin_id"] = a.jellyfin_id; + ad["similarity"] = a.similarity; + ad["bbox"] = nb::make_tuple(a.bbox.x, a.bbox.y, a.bbox.width, a.bbox.height); + actors.append(ad); + } + d["visible_actors"] = actors; + return d; + } + + static SceneAnnotation from_python(nb::object) { + return {}; // never called + } +}; + +// Intermediates: never cross the seam. Provide stubs so register_full_type compiles. +template<> struct PythonConverter { + static constexpr const char* type_name = "TrackedSceneFrame"; + static nb::object to_python(const TrackedSceneFrame&) { return nb::none(); } + static TrackedSceneFrame from_python(nb::object) { return {}; } +}; +template<> struct PythonConverter { + static constexpr const char* type_name = "MatchedSceneFrame"; + static nb::object to_python(const MatchedSceneFrame&) { return nb::none(); } + static MatchedSceneFrame from_python(nb::object) { return {}; } +}; + +} // namespace kpn + +// ── Config from Python dict ───────────────────────────────────────────────────── +// Only the knobs relevant to the replayed chain; everything else keeps its default. + +static Config config_from_dict(nb::dict d) { + Config cfg; + auto getf = [&](const char* k, float& dst) { if (d.contains(k)) dst = nb::cast(d[k]); }; + auto geti = [&](const char* k, int& dst) { if (d.contains(k)) dst = nb::cast(d[k]); }; + auto getd = [&](const char* k, double& dst){ if (d.contains(k)) dst = nb::cast(d[k]); }; + // identity matcher + getf("match_prior", cfg.match_prior); + getf("prob_threshold", cfg.prob_threshold); + getf("match_threshold", cfg.match_threshold); + getf("match_ratio", cfg.match_ratio); + getf("match_ratio_ceil", cfg.match_ratio_ceil); + // face tracker + getf("track_alpha", cfg.track_alpha); + getf("track_min_iou", cfg.track_min_iou); + getf("track_max_embed_dist", cfg.track_max_embed_dist); + geti("track_max_frames_missing", cfg.track_max_frames_missing); + getf("cut_revive_sim", cfg.cut_revive_sim); + geti("cut_inactive_max_frames", cfg.cut_inactive_max_frames); + // scene tracker + getd("extinction_sec", cfg.extinction_sec); + getd("anneal_sec", cfg.anneal_sec); + // gallery expansion (usually off for sweeps; expose so it can be toggled) + if (d.contains("expand_gallery")) cfg.expand_gallery = nb::cast(d["expand_gallery"]); + return cfg; +} + +using Net = kpn::python::PyNetwork; + +NB_MODULE(sae_kpn, m) { + m.doc() = "Real KPN downstream nodes (tracker/matcher/scene_tracker) for Python replay sweeps"; + + kpn::python::register_py_network(m, "Network"); + + // Register converters + channel factories for all four channel types on a net. + // register_py_network doesn't do this (auto_bind does); we patch __init__ to. + // Simpler: expose a free helper the Python side calls right after construction. + m.def("_register_types", [](Net& net) { + net.register_full_type( + [](const EmbeddedSceneFrame& v){ return kpn::PythonConverter::to_python(v); }, + [](nb::object o){ return kpn::PythonConverter::from_python(std::move(o)); }, + "EmbeddedSceneFrame"); + net.register_full_type( + [](const TrackedSceneFrame& v){ return kpn::PythonConverter::to_python(v); }, + [](nb::object o){ return kpn::PythonConverter::from_python(std::move(o)); }, + "TrackedSceneFrame"); + net.register_full_type( + [](const MatchedSceneFrame& v){ return kpn::PythonConverter::to_python(v); }, + [](nb::object o){ return kpn::PythonConverter::from_python(std::move(o)); }, + "MatchedSceneFrame"); + net.register_full_type( + [](const SceneAnnotation& v){ return kpn::PythonConverter::to_python(v); }, + [](nb::object o){ return kpn::PythonConverter::from_python(std::move(o)); }, + "SceneAnnotation"); + }); + + m.def("add_node_python", [](Net& net, std::string name, nb::object callable, + std::vector ins, std::vector outs, + std::size_t cap) { + net.add_node_python(std::move(name), std::move(callable), std::move(ins), + std::move(outs), cap); + }, "net"_a, "name"_a, "callable"_a, "inputs"_a, "outputs"_a, "capacity"_a = 5); + + // ── Real node factories ───────────────────────────────────────────────────── + m.def("add_face_tracker", [](Net& net, std::string name, nb::dict cfg_dict, std::size_t cap) { + Config cfg = config_from_dict(cfg_dict); + auto node = std::make_shared, kpn::out<"tracked">>>(cap, cfg); + net.add(std::move(name), std::move(node)); + }, "net"_a, "name"_a, "config"_a, "capacity"_a = 16); + + m.def("add_identity_matcher", [](Net& net, std::string name, std::string gallery_path, + nb::dict cfg_dict, std::size_t cap) { + Config cfg = config_from_dict(cfg_dict); + cfg.gallery_path = gallery_path; // needed to persist refreshed calibration back + // Cache loaded galleries by path so a threshold sweep (many networks, same + // gallery) pays the ~24s JSON parse only once. The matcher holds a const + // ref; the cache keeps the gallery alive for the process lifetime. + static std::map> cache; + auto it = cache.find(gallery_path); + if (it == cache.end()) + it = cache.emplace(gallery_path, + std::make_shared(load_gallery(gallery_path))).first; + auto node = std::make_shared, kpn::out<"matched">>>( + cap, *it->second, cfg); + net.add(std::move(name), std::move(node)); + }, "net"_a, "name"_a, "gallery"_a, "config"_a, "capacity"_a = 16); + + m.def("add_scene_tracker", [](Net& net, std::string name, nb::dict cfg_dict, std::size_t cap) { + Config cfg = config_from_dict(cfg_dict); + auto node = std::make_shared, kpn::out<"annotation">>>(cap, cfg); + net.add(std::move(name), std::move(node)); + }, "net"_a, "name"_a, "config"_a, "capacity"_a = 16); + + // ── Runtime setters (persistent-pipeline reuse across a threshold sweep) ───── + // Build the network once, then change thresholds between replays — no rebuild, + // no teardown (which is where the ROCm deadlock lives), no gallery reload. + using MatcherWrap = kpn::ObjectVariantNodeWrapper< + IdentityMatcherFunc, SaeVariant, kpn::in<"tracked">, kpn::out<"matched">>; + using SceneWrap = kpn::ObjectVariantNodeWrapper< + SceneTrackerFunc, SaeVariant, kpn::in<"matched">, kpn::out<"annotation">>; + + m.def("set_prob_threshold", [](Net& net, std::string name, float t) { + auto* w = dynamic_cast(net.node_ptr(name)); + if (!w) throw std::runtime_error("set_prob_threshold: '" + name + "' is not an identity_matcher"); + w->functor().set_prob_threshold(t); + }, "net"_a, "name"_a, "value"_a); + + m.def("set_extinction_sec", [](Net& net, std::string name, double s) { + auto* w = dynamic_cast(net.node_ptr(name)); + if (!w) throw std::runtime_error("set_extinction_sec: '" + name + "' is not a scene_tracker"); + w->functor().set_extinction_sec(s); + }, "net"_a, "name"_a, "value"_a); +} diff --git a/src/nodes/camera_position_change_detector_node.hpp b/src/nodes/camera_position_change_detector_node.hpp new file mode 100644 index 0000000..b8b6087 --- /dev/null +++ b/src/nodes/camera_position_change_detector_node.hpp @@ -0,0 +1,73 @@ +#pragma once +#include "types.hpp" +#include "config.hpp" + +#include +#include +#include +#include + +// ── CameraPositionChangeDetectorFunc ────────────────────────────────────────── +// KPN node: flags intra-scene camera-angle changes (hard cuts) by comparing each +// frame's grayscale histogram to the previous frame's. When the normalised +// histogram correlation drops below cut_threshold, the frame is marked with +// Frame::is_cut = true. +// +// This is the pipeline's cut *trigger*: it owns cut detection so a single node +// decides when the camera position has changed, and every downstream stage reads +// the decision off Frame::is_cut (which rides with the frame's timestamp_sec / +// frame_idx). The face_tracker consumes it to re-associate tracks across the cut +// rather than blindly resetting; the track_gallery consumes it to bound +// promotion to a single physical viewpoint. +// +// This is the always-on histogram cut — separate from the opt-in TransNetV2 +// scene detector, which localises true shot boundaries as Frame::is_scene_boundary. +// +// The node is a pure pass-through: it forwards the Frame unchanged except for +// is_cut, so it slots between frame_source and face_detector without altering the +// downstream contract. eof frames are forwarded immediately without processing. + +struct CameraPositionChangeDetectorFunc { + static constexpr std::string_view label() { return "camera_position_change_detector"; } + + explicit CameraPositionChangeDetectorFunc(const Config& cfg) + : cut_threshold_(cfg.cut_threshold) + { + std::cerr << "[camera_position_change_detector] cut_threshold=" + << cut_threshold_ << "\n"; + } + + Frame operator()(Frame f) { + if (f.eof) return f; + + cv::Mat gray; + cv::cvtColor(f.image, 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); + // Cut score in [0,1]: 0 = identical to previous frame, ~1 = fully + // different. Rides on the Frame for the preview HUD / debugging. + f.cut_score = static_cast(std::clamp(1.0 - corr, 0.0, 1.0)); + f.is_cut = (corr < cut_threshold_); + if (f.is_cut) + std::cerr << "[camera_position_change_detector] cut at t=" + << f.timestamp_sec << "s hist_corr=" << corr << "\n"; + } + prev_hist_ = hist; + prev_hist_valid_ = true; + + return f; + } + +private: + float cut_threshold_; + cv::Mat prev_hist_; + bool prev_hist_valid_{false}; +}; diff --git a/src/nodes/embedding_dump_node.hpp b/src/nodes/embedding_dump_node.hpp new file mode 100644 index 0000000..1541dc1 --- /dev/null +++ b/src/nodes/embedding_dump_node.hpp @@ -0,0 +1,124 @@ +#pragma once +#include "types.hpp" +#include "config.hpp" + +#include + +#include +#include +#include +#include + +// ── EmbeddingDumpFunc ───────────────────────────────────────────────────────── +// KPN sink that taps the EmbeddedSceneFrame channel and writes the per-frame face +// metadata + embeddings to one HDF5 file (schema: scripts/optimizer/SCHEMA.md). +// The dump is the expensive, parameter-independent half of the pipeline +// (decode→detect→align→embed); replaying it lets a threshold sweep re-run the cheap +// downstream nodes thousands of times with no GPU. See sae_kpn / scripts/optimizer. +// +// Accumulates in flat/ragged arrays and writes once on EOF. + +struct EmbeddingDumpFunc { + static constexpr std::string_view label() { return "embedding_dump"; } + + EmbeddingDumpFunc(const Config& cfg, std::atomic& done) + : path_(cfg.dump_embeddings_path), movie_(cfg.movie_path), + sample_fps_(cfg.sample_fps), done_(done) + { + std::cerr << "[embedding_dump] writing " << path_ << "\n"; + } + + void operator()(EmbeddedSceneFrame ef) { + if (ef.source.eof) { flush(); return; } + + const int32_t n = static_cast(ef.faces.size()); + ts_.push_back(ef.source.timestamp_sec); + fidx_.push_back(ef.source.frame_idx); + is_cut_.push_back(ef.source.is_cut ? 1 : 0); + is_bnd_.push_back(ef.source.is_scene_boundary ? 1 : 0); + face_off_.push_back(static_cast(conf_.size())); + face_cnt_.push_back(n); + + for (int i = 0; i < n; ++i) { + const auto& f = ef.faces[i]; + bbox_.insert(bbox_.end(), {f.bbox.x, f.bbox.y, f.bbox.width, f.bbox.height}); + for (int k = 0; k < 5; ++k) { + lmk_.push_back(f.landmarks[k].x); + lmk_.push_back(f.landmarks[k].y); + } + conf_.push_back(f.confidence); + const auto& e = ef.embeddings[i]; + emb_.insert(emb_.end(), e.begin(), e.end()); + } + } + + void flush() { + if (written_.exchange(true)) return; + try { + write_hdf5(); + } catch (const H5::Exception& e) { + std::cerr << "[embedding_dump] HDF5 error: " << e.getDetailMsg() << "\n"; + } + done_.store(true, std::memory_order_release); + } + +private: + static constexpr int kSchemaVersion = 1; + static constexpr int kEmbedDim = 512; + + template + void write_vec(H5::Group& g, const char* name, const std::vector& v, + const H5::PredType& dtype, hsize_t cols = 0) { + hsize_t rows = cols ? v.size() / cols : v.size(); + std::vector dims = cols ? std::vector{rows, cols} + : std::vector{rows}; + H5::DataSpace space(static_cast(dims.size()), dims.data()); + auto ds = g.createDataSet(name, dtype, space); + if (!v.empty()) ds.write(v.data(), dtype); + } + + void write_hdf5() { + H5::H5File file(path_, H5F_ACC_TRUNC); + + // root attrs + auto scalar = H5::DataSpace(H5S_SCALAR); + auto ver = file.createAttribute("schema_version", H5::PredType::NATIVE_INT, scalar); + int sv = kSchemaVersion; ver.write(H5::PredType::NATIVE_INT, &sv); + auto ed = file.createAttribute("embed_dim", H5::PredType::NATIVE_INT, scalar); + int dim = kEmbedDim; ed.write(H5::PredType::NATIVE_INT, &dim); + auto fps = file.createAttribute("sample_fps", H5::PredType::NATIVE_FLOAT, scalar); + fps.write(H5::PredType::NATIVE_FLOAT, &sample_fps_); + H5::StrType str(H5::PredType::C_S1, H5T_VARIABLE); + auto mv = file.createAttribute("movie", str, scalar); + mv.write(str, movie_); + + H5::Group frames = file.createGroup("frames"); + write_vec(frames, "timestamp_sec", ts_, H5::PredType::NATIVE_DOUBLE); + write_vec(frames, "frame_idx", fidx_, H5::PredType::NATIVE_INT64); + write_vec(frames, "is_cut", is_cut_, H5::PredType::NATIVE_UINT8); + write_vec(frames, "is_scene_boundary", is_bnd_, H5::PredType::NATIVE_UINT8); + write_vec(frames, "face_offset", face_off_, H5::PredType::NATIVE_INT64); + write_vec(frames, "face_count", face_cnt_, H5::PredType::NATIVE_INT32); + + H5::Group faces = file.createGroup("faces"); + write_vec(faces, "embedding", emb_, H5::PredType::NATIVE_FLOAT, kEmbedDim); + write_vec(faces, "bbox", bbox_, H5::PredType::NATIVE_FLOAT, 4); + write_vec(faces, "landmarks", lmk_, H5::PredType::NATIVE_FLOAT, 10); + write_vec(faces, "confidence", conf_, H5::PredType::NATIVE_FLOAT); + + std::cerr << "[embedding_dump] wrote " << ts_.size() << " frames, " + << conf_.size() << " faces → " << path_ << "\n"; + } + + std::string path_, movie_; + float sample_fps_; + std::atomic& done_; + std::atomic written_{false}; + + std::vector ts_; + std::vector fidx_; + std::vector is_cut_, is_bnd_; + std::vector face_off_; + std::vector face_cnt_; + std::vector emb_, bbox_, lmk_, conf_; +}; diff --git a/src/nodes/scene_detector_node.hpp b/src/nodes/scene_detector_node.hpp new file mode 100644 index 0000000..8ec954a --- /dev/null +++ b/src/nodes/scene_detector_node.hpp @@ -0,0 +1,179 @@ +#pragma once +#include "types.hpp" +#include "config.hpp" +#include "inference/scene_detector.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +// ── SceneDetectorFunc ───────────────────────────────────────────────────────── +// KPN sink node: TransNetV2 shot-boundary detection on the dense frame stream. +// +// Buffers incoming (dense, native-rate) Frames into a rolling window of +// ISceneDetector::kWindow (=100) frames. Every `stride` frames it runs one +// inference and reads back per-frame boundary probabilities, but only trusts the +// central region of each window — TransNetV2 (like most sliding-window boundary +// models) is unreliable near the window edges where it lacks temporal context. +// Overlapping windows by (kWindow - stride) frames means every frame is scored +// from at least one window's trusted centre. +// +// Boundaries (prob > scene_threshold, local maxima) are collected with their +// timestamps and written to scenes.json alongside the main annotations output on +// EOF. This branch is terminal: it produces no pipeline messages, only a file. + +struct SceneDetectorFunc { + static constexpr std::string_view label() { return "scene_detector"; } + + SceneDetectorFunc(const Config& cfg, std::atomic& done) + : detector_(make_scene_detector(cfg)) + , threshold_(cfg.scene_threshold) + , stride_(std::clamp(cfg.scene_stride, 1, ISceneDetector::kWindow)) + , output_path_(scenes_path(cfg.output_path)) + , movie_path_(cfg.movie_path) + , done_(done) + { + // Trusted centre half of each window. Frames outside [guard, kWindow-guard) + // are re-scored by an adjacent window, so we ignore them here to avoid + // edge artefacts and double-counting. + guard_ = (ISceneDetector::kWindow - stride_) / 2; + std::cerr << "[scene_detector] threshold=" << threshold_ + << " stride=" << stride_ + << " guard=" << guard_ + << " output=" << output_path_ << "\n"; + } + + void operator()(Frame f) { + if (f.eof) { + flush_remaining(); + write_output(); + done_.store(true, std::memory_order_release); + return; + } + + images_.push_back(f.image); + times_.push_back(f.timestamp_sec); + + // Once we have a full window, score it and slide forward by `stride`. + while (static_cast(images_.size()) >= ISceneDetector::kWindow) { + score_window(); + for (int i = 0; i < stride_; ++i) { + images_.pop_front(); + times_.pop_front(); + } + window_base_ += stride_; + } + } + +private: + // Run TransNetV2 on the leading kWindow frames of the buffer and record any + // boundaries found within the trusted centre region. + void score_window() { + std::vector win(images_.begin(), + images_.begin() + ISceneDetector::kWindow); + std::vector probs = detector_->detect_window(win); + + // On the very first window there is no preceding window, so trust from 0; + // otherwise skip the leading guard already covered by the previous window. + const int lo = (window_base_ == 0) ? 0 : guard_; + const int hi = ISceneDetector::kWindow - guard_; + for (int i = lo; i < hi; ++i) { + if (probs[i] <= threshold_) continue; + // Local maximum → the boundary frame (avoid a run of high scores + // registering as several adjacent cuts). + const bool peak = + (i == 0 || probs[i] >= probs[i-1]) && + (i == kLast_() || probs[i] >= probs[i+1]); + if (peak) + boundaries_.push_back({times_[i], probs[i]}); + } + } + + // At EOF the tail (< kWindow frames) never formed a full window. Pad it out + // to kWindow by repeating the last frame so the final real frames still get + // scored, then take only the region past what earlier windows covered. + void flush_remaining() { + const int n = static_cast(images_.size()); + if (n == 0) return; + std::vector win(images_.begin(), images_.end()); + cv::Mat last = win.back(); + while (static_cast(win.size()) < ISceneDetector::kWindow) + win.push_back(last); + + std::vector probs = detector_->detect_window(win); + const int lo = (window_base_ == 0) ? 0 : guard_; + for (int i = lo; i < n; ++i) { // only real (non-padded) frames + if (probs[i] <= threshold_) continue; + const bool peak = + (i == 0 || probs[i] >= probs[i-1]) && + (i == n - 1 || probs[i] >= probs[i+1]); + if (peak) + boundaries_.push_back({times_[i], probs[i]}); + } + } + + void write_output() { + if (written_) return; + written_ = true; + + // Merge boundaries closer than one frame apart (dedup across window seams). + std::sort(boundaries_.begin(), boundaries_.end(), + [](const Boundary& a, const Boundary& b) { + return a.t < b.t; + }); + + nlohmann::json root; + root["schema_version"] = 1; + root["movie"] = movie_path_; + root["model"] = "transnetv2"; + root["threshold"] = threshold_; + nlohmann::json cuts = nlohmann::json::array(); + double last_t = -1e9; + for (const auto& b : boundaries_) { + if (b.t - last_t < 0.04) continue; // ~1 frame @25fps dedup + cuts.push_back({{"t", b.t}, {"probability", b.prob}}); + last_t = b.t; + } + root["cuts"] = std::move(cuts); + + std::ofstream f(output_path_); + if (!f.is_open()) { + std::cerr << "\n[scene_detector] ERROR: cannot write " + << output_path_ << "\n"; + return; + } + f << root.dump(2) << "\n"; + std::cerr << "\n[scene_detector] wrote " << root["cuts"].size() + << " boundaries → " << output_path_ << "\n"; + } + + static int kLast_() { return ISceneDetector::kWindow - 1; } + + // annotations.json → annotations.scenes.json (or scenes.json for bare names) + static std::string scenes_path(const std::string& out) { + auto dot = out.find_last_of('.'); + if (dot == std::string::npos) return out + ".scenes.json"; + return out.substr(0, dot) + ".scenes.json"; + } + + struct Boundary { double t; float prob; }; + + std::unique_ptr detector_; + float threshold_; + int stride_; + int guard_{0}; + std::string output_path_; + std::string movie_path_; + + std::atomic& done_; + std::deque images_; + std::deque times_; + int64_t window_base_{0}; // frame index of images_.front() + std::vector boundaries_; + bool written_{false}; +};