Initial commit: scene-actor-extraction pipeline
Source (KPN++ pipeline nodes, ArcFace embedders, SCRFD/YuNet detectors, gallery builder), build scripts, and eval artifacts. - external/KPN as a git submodule (gitea.tourolle.paris/dtourolle/KPN) - ONNX models tracked via Git LFS (models/*.onnx) - generated outputs, TensorRT engines, reference repos, and media ignored
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
// scene_preview — same pipeline as scene_analyze but with a live annotated
|
||||
// display window driven from the main thread.
|
||||
//
|
||||
// KPN topology:
|
||||
//
|
||||
// [frame_source] ──► [face_detector] ──► [face_aligner] ──► [embedder]
|
||||
// ──► [identity_matcher] ──► FanoutNode<MatchedSceneFrame,2>
|
||||
// ├──► [scene_tracker] ──► [result_sink] (background thread)
|
||||
// └──► [preview_node] (main thread)
|
||||
//
|
||||
// The main thread drives preview_node via preview.step(). When the movie ends
|
||||
// (MatchedSceneFrame.source.eof == true) operator() returns false, the loop
|
||||
// ends, and net.stop() is called. result_sink writes its JSON before its
|
||||
// thread is joined by net.stop(), so the output file is always complete.
|
||||
//
|
||||
// Usage: same flags as scene_analyze (see main.cpp for reference).
|
||||
// --preview-width <px> max display width in pixels (default: 1280)
|
||||
|
||||
#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"
|
||||
#include "nodes/preview_node.hpp"
|
||||
|
||||
#include <kpn/kpn.hpp>
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
|
||||
static Config parse_args(int argc, char** argv) {
|
||||
Config cfg;
|
||||
cfg.detector_model = kDefaultDetectorModel;
|
||||
cfg.arcface_model = kDefaultArcfaceModel;
|
||||
cfg.output_path = "annotations.json";
|
||||
cfg.verbosity = Verbosity::standard; // default to standard in preview mode
|
||||
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
auto arg = [&](const char* f) { return std::strcmp(argv[i], f) == 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());
|
||||
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;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
ActorGallery gallery;
|
||||
try { gallery = load_gallery(cfg.gallery_path); }
|
||||
catch (const std::exception& e) {
|
||||
std::cerr << "Gallery error: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
// ── 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};
|
||||
|
||||
// ── KPN ObjectNodes ───────────────────────────────────────────────────────
|
||||
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);
|
||||
|
||||
// MainThreadNode — no thread spawned; driven by preview.step() below
|
||||
PreviewNode preview{cfg, 16};
|
||||
|
||||
// matcher → FanoutNode<MatchedSceneFrame,2> → [scene_tracker, preview] (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">(), preview.input<"matched">()),
|
||||
kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">())
|
||||
);
|
||||
|
||||
// ── Run ───────────────────────────────────────────────────────────────────
|
||||
std::cerr << "[main] starting pipeline — press q or Esc to quit early\n";
|
||||
net.start();
|
||||
|
||||
// Main thread drives the display window; returns false on EOF or q/Esc
|
||||
while (preview.step()) {
|
||||
cv::waitKey(1); // pump OS events between frames
|
||||
}
|
||||
|
||||
net.stop();
|
||||
sink_fn.flush(); // write whatever was accumulated (no-op if EOF already flushed)
|
||||
cv::destroyAllWindows();
|
||||
net.print_diagnostics();
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user