feat(engine): add Python replay bindings, gallery pose-expansion, scene detection, embedding dumps
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.
This commit is contained in:
@@ -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 <path> --out <dump.h5> [--arcface <model.onnx>]
|
||||
// [--detector <model.onnx>] [--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 <kpn/kpn.hpp>
|
||||
|
||||
#include <atomic>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
|
||||
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 <path> --out <dump.h5> "
|
||||
"[--arcface <model.onnx>] [--fps 1] ...\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::atomic<bool> 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<FrameSourceFunc, kpn::in<>, kpn::out<"raw">, "frame_source", 0> source (source_fn, 32);
|
||||
kpn::ObjectNode<CameraPositionChangeDetectorFunc, kpn::in<"raw">, kpn::out<"frame">, "camera_pos", 0> campos (campos_fn, 32);
|
||||
kpn::ObjectNode<FaceDetectorFunc, kpn::in<"frame">, kpn::out<"scene">, "face_detector", 0> detector(detector_fn, 64);
|
||||
kpn::ObjectNode<FaceAlignerFunc, kpn::in<"scene">, kpn::out<"aligned">, "face_aligner", 0> aligner (aligner_fn, 64);
|
||||
kpn::ObjectNode<EmbedderFunc, kpn::in<"aligned">, kpn::out<"embedded">, "embedder", 0> embedder(embedder_fn, 32);
|
||||
kpn::ObjectNode<EmbeddingDumpFunc, kpn::in<"embedded">,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;
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <iostream>
|
||||
#include <limits>
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#ifdef SAE_DEBUG
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#endif
|
||||
#include <opencv2/core.hpp>
|
||||
#include <filesystem>
|
||||
|
||||
// ── 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<AnnexEntry>& 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<BufEntry> buf;
|
||||
std::map<int, int> 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<int>(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<int>(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<BufEntry>& buf) {
|
||||
float min_sim = std::numeric_limits<float>::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<float>::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<int, TrackState> tracks_;
|
||||
std::vector<AnnexEntry> annex_;
|
||||
};
|
||||
@@ -0,0 +1,45 @@
|
||||
#pragma once
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include <opencv2/core.hpp>
|
||||
|
||||
// ── 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<float> detect_window(
|
||||
const std::vector<cv::Mat>& 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<ISceneDetector> make_scene_detector(const Config& cfg);
|
||||
@@ -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 <kpn/python/bindings.hpp>
|
||||
#include <kpn/python/object_variant_node.hpp>
|
||||
|
||||
#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 <nanobind/nanobind.h>
|
||||
#include <nanobind/ndarray.h>
|
||||
#include <nanobind/stl/string.h>
|
||||
#include <nanobind/stl/vector.h>
|
||||
#include <nanobind/stl/map.h>
|
||||
|
||||
#include <memory>
|
||||
#include <variant>
|
||||
|
||||
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<EmbeddedSceneFrame, TrackedSceneFrame,
|
||||
MatchedSceneFrame, SceneAnnotation>;
|
||||
|
||||
// ── 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<EmbeddedSceneFrame> {
|
||||
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<nb::dict>(o);
|
||||
EmbeddedSceneFrame ef;
|
||||
ef.source.timestamp_sec = nb::cast<double>(d["timestamp_sec"]);
|
||||
ef.source.frame_idx = d.contains("frame_idx") ? nb::cast<int64_t>(d["frame_idx"]) : -1;
|
||||
ef.source.eof = d.contains("eof") ? nb::cast<bool>(d["eof"]) : false;
|
||||
ef.source.is_cut = d.contains("is_cut") ? nb::cast<bool>(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::ndarray<float, nb::shape<-1, 4>, nb::c_contig>>(d["bbox"]);
|
||||
auto lmk = nb::cast<nb::ndarray<float, nb::shape<-1, 10>, nb::c_contig>>(d["landmarks"]);
|
||||
auto conf = nb::cast<nb::ndarray<float, nb::shape<-1>, nb::c_contig>>(d["confidence"]);
|
||||
auto emb = nb::cast<nb::ndarray<float, nb::shape<-1, 512>, 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<SceneAnnotation> {
|
||||
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<TrackedSceneFrame> {
|
||||
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<MatchedSceneFrame> {
|
||||
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<float>(d[k]); };
|
||||
auto geti = [&](const char* k, int& dst) { if (d.contains(k)) dst = nb::cast<int>(d[k]); };
|
||||
auto getd = [&](const char* k, double& dst){ if (d.contains(k)) dst = nb::cast<double>(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<bool>(d["expand_gallery"]);
|
||||
return cfg;
|
||||
}
|
||||
|
||||
using Net = kpn::python::PyNetwork<SaeVariant>;
|
||||
|
||||
NB_MODULE(sae_kpn, m) {
|
||||
m.doc() = "Real KPN downstream nodes (tracker/matcher/scene_tracker) for Python replay sweeps";
|
||||
|
||||
kpn::python::register_py_network<SaeVariant>(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<EmbeddedSceneFrame>(
|
||||
[](const EmbeddedSceneFrame& v){ return kpn::PythonConverter<EmbeddedSceneFrame>::to_python(v); },
|
||||
[](nb::object o){ return kpn::PythonConverter<EmbeddedSceneFrame>::from_python(std::move(o)); },
|
||||
"EmbeddedSceneFrame");
|
||||
net.register_full_type<TrackedSceneFrame>(
|
||||
[](const TrackedSceneFrame& v){ return kpn::PythonConverter<TrackedSceneFrame>::to_python(v); },
|
||||
[](nb::object o){ return kpn::PythonConverter<TrackedSceneFrame>::from_python(std::move(o)); },
|
||||
"TrackedSceneFrame");
|
||||
net.register_full_type<MatchedSceneFrame>(
|
||||
[](const MatchedSceneFrame& v){ return kpn::PythonConverter<MatchedSceneFrame>::to_python(v); },
|
||||
[](nb::object o){ return kpn::PythonConverter<MatchedSceneFrame>::from_python(std::move(o)); },
|
||||
"MatchedSceneFrame");
|
||||
net.register_full_type<SceneAnnotation>(
|
||||
[](const SceneAnnotation& v){ return kpn::PythonConverter<SceneAnnotation>::to_python(v); },
|
||||
[](nb::object o){ return kpn::PythonConverter<SceneAnnotation>::from_python(std::move(o)); },
|
||||
"SceneAnnotation");
|
||||
});
|
||||
|
||||
m.def("add_node_python", [](Net& net, std::string name, nb::object callable,
|
||||
std::vector<std::string> ins, std::vector<std::string> 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::ObjectVariantNodeWrapper<
|
||||
FaceTrackerFunc, SaeVariant, kpn::in<"embedded">, 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<std::string, std::shared_ptr<ActorGallery>> cache;
|
||||
auto it = cache.find(gallery_path);
|
||||
if (it == cache.end())
|
||||
it = cache.emplace(gallery_path,
|
||||
std::make_shared<ActorGallery>(load_gallery(gallery_path))).first;
|
||||
auto node = std::make_shared<kpn::ObjectVariantNodeWrapper<
|
||||
IdentityMatcherFunc, SaeVariant, kpn::in<"tracked">, 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::ObjectVariantNodeWrapper<
|
||||
SceneTrackerFunc, SaeVariant, kpn::in<"matched">, 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<MatcherWrap*>(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<SceneWrap*>(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);
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <algorithm>
|
||||
#include <iostream>
|
||||
#include <string_view>
|
||||
|
||||
// ── 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<float>(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};
|
||||
};
|
||||
@@ -0,0 +1,124 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <H5Cpp.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// ── 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<bool>& 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<int32_t>(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<int64_t>(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<typename T>
|
||||
void write_vec(H5::Group& g, const char* name, const std::vector<T>& v,
|
||||
const H5::PredType& dtype, hsize_t cols = 0) {
|
||||
hsize_t rows = cols ? v.size() / cols : v.size();
|
||||
std::vector<hsize_t> dims = cols ? std::vector<hsize_t>{rows, cols}
|
||||
: std::vector<hsize_t>{rows};
|
||||
H5::DataSpace space(static_cast<int>(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<bool>& done_;
|
||||
std::atomic<bool> written_{false};
|
||||
|
||||
std::vector<double> ts_;
|
||||
std::vector<int64_t> fidx_;
|
||||
std::vector<uint8_t> is_cut_, is_bnd_;
|
||||
std::vector<int64_t> face_off_;
|
||||
std::vector<int32_t> face_cnt_;
|
||||
std::vector<float> emb_, bbox_, lmk_, conf_;
|
||||
};
|
||||
@@ -0,0 +1,179 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
#include "inference/scene_detector.hpp"
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <deque>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// ── 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<bool>& 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<int>(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<cv::Mat> win(images_.begin(),
|
||||
images_.begin() + ISceneDetector::kWindow);
|
||||
std::vector<float> 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<int>(images_.size());
|
||||
if (n == 0) return;
|
||||
std::vector<cv::Mat> win(images_.begin(), images_.end());
|
||||
cv::Mat last = win.back();
|
||||
while (static_cast<int>(win.size()) < ISceneDetector::kWindow)
|
||||
win.push_back(last);
|
||||
|
||||
std::vector<float> 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<ISceneDetector> detector_;
|
||||
float threshold_;
|
||||
int stride_;
|
||||
int guard_{0};
|
||||
std::string output_path_;
|
||||
std::string movie_path_;
|
||||
|
||||
std::atomic<bool>& done_;
|
||||
std::deque<cv::Mat> images_;
|
||||
std::deque<double> times_;
|
||||
int64_t window_base_{0}; // frame index of images_.front()
|
||||
std::vector<Boundary> boundaries_;
|
||||
bool written_{false};
|
||||
};
|
||||
Reference in New Issue
Block a user