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:
+213
@@ -0,0 +1,213 @@
|
||||
// scene_analyze — identify actors in a movie using a KPN pipeline
|
||||
//
|
||||
// KPN topology (release build):
|
||||
//
|
||||
// [frame_source] ──Frame──► [face_detector] ──SceneFrame──► [face_aligner]
|
||||
// ──AlignedSceneFrame──► [embedder] ──EmbeddedSceneFrame──►
|
||||
// [identity_matcher] ──MatchedSceneFrame──► [scene_tracker]
|
||||
// ──SceneAnnotation──► [result_sink]
|
||||
//
|
||||
// Debug build (SAE_DEBUG=1):
|
||||
// [identity_matcher] output fans out to both [scene_tracker] AND [debug_renderer].
|
||||
// FanoutNode<MatchedSceneFrame, 2> is auto-inserted by make_network().
|
||||
//
|
||||
// Usage:
|
||||
// scene_analyze --movie <path> --gallery <gallery.json> [options]
|
||||
//
|
||||
// Options:
|
||||
// --output <path> output JSON (default: annotations.json)
|
||||
// --fps <N> sample rate in frames/sec (default: 1.0)
|
||||
// --verbosity <0|1|2> 0=minimal, 1=standard, 2=jellyfin-xray (default: 0)
|
||||
// --match-threshold <f> cosine dist threshold (default: 0.45)
|
||||
// --extinction <f> actor extinction window in seconds (default: 5.0)
|
||||
// --detector <path> override SCRFD detector model path
|
||||
// --arcface <path> override ArcFace model path
|
||||
// --max-faces <N> max faces kept per frame (default: 10)
|
||||
// (SAE_DEBUG only)
|
||||
// --debug-dir <path> debug frames output dir (default: debug_frames)
|
||||
// --crop-context <f> bbox expansion factor for context crops (default: 1.5)
|
||||
|
||||
#include "config.hpp"
|
||||
#include "ort_provider.hpp"
|
||||
#include "types.hpp"
|
||||
#include "gallery/gallery_store.hpp"
|
||||
#include "nodes/frame_source_node.hpp"
|
||||
#include "nodes/face_detector_node.hpp"
|
||||
#include "nodes/face_aligner_node.hpp"
|
||||
#include "nodes/embedder_node.hpp"
|
||||
#include "nodes/face_tracker_node.hpp"
|
||||
#include "nodes/identity_matcher_node.hpp"
|
||||
#include "nodes/scene_tracker_node.hpp"
|
||||
#include "nodes/result_sink_node.hpp"
|
||||
#ifdef SAE_DEBUG
|
||||
#include "nodes/debug_renderer_node.hpp"
|
||||
#endif
|
||||
|
||||
#include <kpn/kpn.hpp>
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
|
||||
// ── CLI parsing ───────────────────────────────────────────────────────────────
|
||||
|
||||
static Config parse_args(int argc, char** argv) {
|
||||
Config cfg;
|
||||
cfg.detector_model = kDefaultDetectorModel;
|
||||
cfg.arcface_model = kDefaultArcfaceModel;
|
||||
cfg.output_path = "annotations.json";
|
||||
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
auto arg = [&](const char* flag) { return std::strcmp(argv[i], flag) == 0; };
|
||||
auto next = [&]() -> std::string {
|
||||
if (++i >= argc) throw std::runtime_error(std::string("missing arg after ") + argv[i-1]);
|
||||
return argv[i];
|
||||
};
|
||||
|
||||
if (arg("--movie")) cfg.movie_path = next();
|
||||
else if (arg("--gallery")) cfg.gallery_path = next();
|
||||
else if (arg("--output")) cfg.output_path = next();
|
||||
else if (arg("--fps")) cfg.sample_fps = std::stof(next());
|
||||
else if (arg("--max-decode-fps")) cfg.max_decode_fps = std::stof(next());
|
||||
else if (arg("--start")) cfg.start_sec = std::stod(next());
|
||||
else if (arg("--end")) cfg.end_sec = std::stod(next());
|
||||
else if (arg("--cut-threshold")) cfg.cut_threshold = std::stof(next());
|
||||
else if (arg("--verbosity")) { int v = std::stoi(next()); cfg.verbosity = v == 2 ? Verbosity::xray : v == 1 ? Verbosity::standard : Verbosity::minimal; }
|
||||
else if (arg("--prior")) cfg.match_prior = std::stof(next());
|
||||
else if (arg("--prob-threshold")) cfg.prob_threshold = std::stof(next());
|
||||
else if (arg("--match-threshold")) cfg.match_threshold = std::stof(next());
|
||||
else if (arg("--extinction")) cfg.extinction_sec = std::stod(next());
|
||||
else if (arg("--detector")) cfg.detector_model = next();
|
||||
else if (arg("--detector-engine")) cfg.detector_engine = next();
|
||||
else if (arg("--arcface")) cfg.arcface_model = next();
|
||||
else if (arg("--arcface-engine")) cfg.arcface_engine = next();
|
||||
else if (arg("--conf")) cfg.detector_conf = std::stof(next());
|
||||
else if (arg("--max-faces")) cfg.max_faces = std::stoi(next());
|
||||
else if (arg("--min-face-px")) cfg.min_face_px = std::stof(next());
|
||||
else if (arg("--ratio")) cfg.match_ratio = std::stof(next());
|
||||
else if (arg("--ratio-ceil")) cfg.match_ratio_ceil = std::stof(next());
|
||||
else if (arg("--track-alpha")) cfg.track_alpha = std::stof(next());
|
||||
else if (arg("--track-min-iou")) cfg.track_min_iou = std::stof(next());
|
||||
else if (arg("--track-max-embed")) cfg.track_max_embed_dist = std::stof(next());
|
||||
else if (arg("--track-max-missing")) cfg.track_max_frames_missing = std::stoi(next());
|
||||
else if (arg("--track-min-frames")) cfg.track_min_frames = std::stoi(next());
|
||||
else if (arg("--anneal")) cfg.anneal_sec = std::stod(next());
|
||||
else if (arg("--trt-cache")) cfg.trt.cache_dir = next();
|
||||
else if (arg("--trt-fp16")) cfg.trt.fp16 = true;
|
||||
else if (arg("--no-trt-fp16")) cfg.trt.fp16 = false;
|
||||
else if (arg("--trt-int8")) cfg.trt.int8 = true;
|
||||
else if (arg("--embed-batch")) cfg.embed_batch_size = std::stoi(next());
|
||||
#ifdef SAE_DEBUG
|
||||
else if (arg("--debug-dir")) cfg.debug_dir = next();
|
||||
else if (arg("--crop-context")) cfg.crop_context = std::stof(next());
|
||||
#endif
|
||||
else {
|
||||
std::cerr << "[warn] unknown flag: " << argv[i] << "\n";
|
||||
}
|
||||
}
|
||||
|
||||
if (cfg.movie_path.empty()) throw std::runtime_error("--movie is required");
|
||||
if (cfg.gallery_path.empty()) throw std::runtime_error("--gallery is required");
|
||||
return cfg;
|
||||
}
|
||||
|
||||
// ── Main ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
Config cfg;
|
||||
try {
|
||||
cfg = parse_args(argc, argv);
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "Usage error: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Load actor gallery
|
||||
ActorGallery gallery;
|
||||
try {
|
||||
gallery = load_gallery(cfg.gallery_path);
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "Gallery error: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
std::cerr << "[main] gallery loaded: " << gallery.actors.size() << " actors\n";
|
||||
|
||||
// ── Construct node functors ───────────────────────────────────────────────
|
||||
|
||||
std::atomic<bool> done{false};
|
||||
|
||||
const OrtProvider provider = detect_ort_provider();
|
||||
std::cerr << "[main] inference provider: " << provider_name(provider) << "\n";
|
||||
|
||||
FrameSourceFunc source_fn {cfg};
|
||||
FaceDetectorFunc detector_fn{cfg, provider};
|
||||
FaceAlignerFunc aligner_fn;
|
||||
EmbedderFunc embedder_fn{cfg, provider};
|
||||
FaceTrackerFunc ftracker_fn{cfg};
|
||||
IdentityMatcherFunc matcher_fn {gallery, cfg};
|
||||
SceneTrackerFunc tracker_fn {cfg};
|
||||
ResultSinkFunc sink_fn {cfg, done};
|
||||
#ifdef SAE_DEBUG
|
||||
DebugRendererFunc debug_fn {cfg};
|
||||
#endif
|
||||
|
||||
// ── Wrap in KPN ObjectNodes ───────────────────────────────────────────────
|
||||
|
||||
// Queue sizes tuned to the pipeline's speed profile:
|
||||
// embedder (16ms) is the slowest GPU node — buffer before it must be largest
|
||||
// to prevent face_aligner pool overflows and frame drops.
|
||||
kpn::ObjectNode<FrameSourceFunc, kpn::in<>, kpn::out<"frame">, "frame_source", 0> source (source_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<FaceTrackerFunc, kpn::in<"embedded">, kpn::out<"tracked">, "face_tracker", 0> ftracker (ftracker_fn, 16);
|
||||
kpn::ObjectNode<IdentityMatcherFunc, kpn::in<"tracked">, kpn::out<"matched">, "identity_matcher", 0> matcher (matcher_fn, 16);
|
||||
kpn::ObjectNode<SceneTrackerFunc, kpn::in<"matched">, kpn::out<"annotation">, "scene_tracker", 0> tracker (tracker_fn, 16);
|
||||
kpn::ObjectNode<ResultSinkFunc, kpn::in<"annotation">,kpn::out<>, "result_sink", 0> sink (sink_fn, 16);
|
||||
|
||||
// ── Build static network ──────────────────────────────────────────────────
|
||||
|
||||
#ifdef SAE_DEBUG
|
||||
kpn::ObjectNode<DebugRendererFunc, kpn::in<"matched">, kpn::out<>, "debug_renderer", 1> debug_node(debug_fn, 16);
|
||||
|
||||
// matcher → FanoutNode<MatchedSceneFrame,2> → scene_tracker + debug_node (auto-inserted)
|
||||
auto net = kpn::make_network(
|
||||
kpn::edge(source.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">(), ftracker.input<"embedded">()),
|
||||
kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()),
|
||||
kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()),
|
||||
kpn::edge(matcher.output<"matched">(), debug_node.input<"matched">()),
|
||||
kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">())
|
||||
);
|
||||
#else
|
||||
auto net = kpn::make_network(
|
||||
kpn::edge(source.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">(), ftracker.input<"embedded">()),
|
||||
kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()),
|
||||
kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()),
|
||||
kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">())
|
||||
);
|
||||
#endif
|
||||
|
||||
// ── Run ───────────────────────────────────────────────────────────────────
|
||||
// Note: StaticNetwork does not expose set_error_handler; node exceptions
|
||||
// are printed to stderr by the KPN run_loop and the network continues.
|
||||
std::cerr << "[main] starting pipeline…\n";
|
||||
net.start();
|
||||
|
||||
// Main thread waits until ResultSinkFunc signals EOF completion
|
||||
while (!done.load(std::memory_order_acquire))
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||
|
||||
net.stop();
|
||||
net.print_diagnostics();
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user