scene-actor-extraction/src/scene_preview.cpp

// scene_preview — same pipeline as scene_analyze but with a live annotated
// display window driven from the main thread.
//
// KPN topology:
//
//   [frame_source] ──► [camera_pos] ──► [face_detector] ──► [face_aligner] ──► [embedder]
//          ──► [identity_matcher] ──► FanoutNode<MatchedSceneFrame,2>
//
// camera_pos (histogram cut detector) stamps Frame::cut_score / is_cut, which
// ride through to the preview HUD's cut-score meter.
//                   ├──► [frame_annotation] ──► [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 "types.hpp"
#include "gallery/embedder_stamp.hpp"
#include "gallery/gallery_store.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/face_tracker_node.hpp"
#include "nodes/identity_matcher_node.hpp"
#include "track_registry.hpp"
#include "evidence_discount.hpp"
#include "nodes/frame_annotation_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 <map>
#include <mutex>
#include <stdexcept>
#include <string>
#include <string_view>
#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("--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("--require-gallery-stamp")) cfg.require_gallery_stamp = true;
        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("--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-min-prob"))      cfg.track_assoc_min_prob   = std::stof(next());
        else if (arg("--track-extinction"))    cfg.track_extinction_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());
        // Per-film gallery expansion — preview supports it (same cfg fields).
        else if (arg("--expand-gallery"))      cfg.expand_gallery         = true;
        else if (arg("--expand-buffer"))       cfg.expand_buffer_size     = std::stoi(next());
        else if (arg("--expand-band-lo"))      cfg.expand_band_lo         = std::stof(next());
        else if (arg("--expand-band-hi"))      cfg.expand_band_hi         = std::stof(next());
        else if (arg("--expand-min-anchor"))   cfg.expand_min_anchor_frames = std::stoi(next());
        // Scene detection is scene_analyze-only (needs the dense TransNetV2 branch).
        // Accept the flags so a shared command line runs, but note they're inert
        // here — the preview shows the histogram cut-score meter instead.
        else if (arg("--scene-detect")) {
            std::cerr << "[preview] note: --scene-detect is inert in preview "
                         "(TransNetV2 needs the dense scene_analyze pipeline); "
                         "showing the histogram cut-score meter instead\n";
        }
        else if (arg("--scene-detector") || arg("--scene-detector-engine") ||
                 arg("--scene-threshold") || arg("--scene-stride") ||
                 arg("--scene-decode-fps") || arg("--dense-scale")) {
            next();  // consume the value; inert in preview
        }
        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);
        /// TRACES: GR-004 | SR-001
        verify_gallery_embedder(gallery, cfg.gallery_path, cfg.arcface_model,
                                cfg.require_gallery_stamp);
    }
    catch (const std::exception& e) {
        std::cerr << "Gallery error: " << e.what() << "\n";
        return 1;
    }

    // ── Functors ──────────────────────────────────────────────────────────────
    std::atomic<bool> done{false};

    FrameSourceFunc                   source_fn  {cfg};
    CameraPositionChangeDetectorFunc  campos_fn  {cfg};
    FaceDetectorFunc    detector_fn{cfg};
    FaceAlignerFunc     aligner_fn;
    EmbedderFunc        embedder_fn{cfg};
    /// TRACES: DP-001, AR-007, AR-012, AR-024 | PR-004, SR-002
    // Construction order matters and is the same as main.cpp's, deliberately:
    // the matcher fits (or loads) the calibration, the registry needs a
    // discounter built from it, and the tracker needs both. DP-001 says modes
    // are front-ends that must not fork pipeline logic -- this file had forked
    // it and then rotted, constructing FaceTrackerFunc{cfg} against a signature
    // that stopped existing with the AR-007/AR-008 redesign, so scene_preview
    // has not compiled since. Keeping the order identical is what stops that
    // recurring.
    IdentityMatcherFunc matcher_fn {gallery, cfg};
    auto same_person = same_person_probability(matcher_fn.calibration());
    TrackRegistry::Config reg_cfg;
    reg_cfg.track_extinction_sec = cfg.track_extinction_sec;
    auto registry = std::make_shared<TrackRegistry>(
        reg_cfg, EvidenceDiscounter(same_person));
    matcher_fn.set_registry(registry);

    FaceTrackerFunc     ftracker_fn{cfg, registry, same_person};
    FrameAnnotationFunc tracker_fn {};
    ResultSinkFunc      sink_fn    {cfg, done};

    // AR-012/AR-016: windows come from registry claims, and tracks still live
    // at EOF must be flushed or the closing scene's cast is never emitted.
    registry->on_track_dead([&sink_fn](const DeadTrack& d) { sink_fn.add_claim(d); });
    sink_fn.set_pre_write_hook([registry](double last_ts) { registry->flush(last_ts); });

    // ── KPN ObjectNodes ───────────────────────────────────────────────────────
    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<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<FrameAnnotationFunc, kpn::in<"matched">,   kpn::out<"annotation">, "frame_annotation", 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> → [frame_annotation, preview]  (auto-inserted)
    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">(),   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">())
    );

    // ── Pipeline observability (KPN event handler) ────────────────────────────
    // Tally dropped frames per node (channel overflow) and detect a node that
    // stops unexpectedly so the preview loop below can bail out instead of
    // spinning on a dead pipeline.
    std::mutex                  event_mtx;
    std::map<std::string, long> overflow_counts;
    std::atomic<bool>           node_crashed{false};

    net.set_event_handler(
        [&](std::string_view node_name, kpn::NodeEvent ev,
            std::chrono::steady_clock::time_point) {
            if (ev == kpn::NodeEvent::Overflow) {
                std::lock_guard<std::mutex> lk(event_mtx);
                ++overflow_counts[std::string(node_name)];
            } else {  // NodeEvent::Closed
                if (node_name == "result_sink" && done.load(std::memory_order_acquire))
                    return;
                std::cerr << "[main] node '" << node_name
                          << "' stopped unexpectedly — aborting pipeline\n";
                node_crashed.store(true, std::memory_order_release);
            }
        });

    // ── 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.
    // Bail out early if a node crashes.
    while (!node_crashed.load(std::memory_order_acquire) && 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();

    {
        std::lock_guard<std::mutex> lk(event_mtx);
        if (!overflow_counts.empty()) {
            std::cerr << "[main] dropped frames (channel overflow):\n";
            for (const auto& [name, count] : overflow_counts)
                std::cerr << "  " << name << ": " << count << "\n";
        }
    }

    return node_crashed.load(std::memory_order_acquire) ? 1 : 0;
}