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,204 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include "config.hpp"
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
using json = nlohmann::json;
|
||||
|
||||
// ── ResultSinkFunc ────────────────────────────────────────────────────────────
|
||||
// KPN sink node: accumulates SceneAnnotations and writes the final JSON on EOF.
|
||||
//
|
||||
// Verbosity::minimal — merges per-frame presence into contiguous time windows.
|
||||
// Output: { "movie": "...", "actors": [{ "name", "imdb_id", "scenes": [[t0,t1], ...] }] }
|
||||
//
|
||||
// Verbosity::standard — per-frame detail including bboxes, similarity, unknowns.
|
||||
// Output: { "frames": [{ "t", "identified": [...], "unknowns": [...] }] }
|
||||
//
|
||||
// eof signal: sets done_ = true so the main thread can call net.stop().
|
||||
|
||||
struct ResultSinkFunc {
|
||||
static constexpr std::string_view label() { return "result_sink"; }
|
||||
|
||||
ResultSinkFunc(const Config& cfg, std::atomic<bool>& done)
|
||||
: cfg_(cfg), done_(done)
|
||||
{}
|
||||
|
||||
void operator()(SceneAnnotation sa) {
|
||||
if (sa.eof) {
|
||||
flush();
|
||||
return;
|
||||
}
|
||||
|
||||
// Progress to stderr
|
||||
std::cerr << "\r[result_sink] t=" << sa.timestamp_sec << "s"
|
||||
<< " active=" << count_known(sa.visible_actors)
|
||||
<< " unknowns=" << count_unknown(sa.visible_actors)
|
||||
<< std::flush;
|
||||
|
||||
frames_.push_back(std::move(sa));
|
||||
}
|
||||
|
||||
// Write accumulated results and signal done. Safe to call more than once.
|
||||
void flush() {
|
||||
if (written_.exchange(true)) return;
|
||||
write_output();
|
||||
done_.store(true, std::memory_order_release);
|
||||
}
|
||||
|
||||
private:
|
||||
static int count_known(const std::vector<IdentifiedActor>& v) {
|
||||
int n = 0;
|
||||
for (const auto& a : v) if (a.actor_idx >= 0) ++n;
|
||||
return n;
|
||||
}
|
||||
static int count_unknown(const std::vector<IdentifiedActor>& v) {
|
||||
int n = 0;
|
||||
for (const auto& a : v) if (a.actor_idx < 0) ++n;
|
||||
return n;
|
||||
}
|
||||
|
||||
void write_output() {
|
||||
std::cerr << "\n[result_sink] writing " << cfg_.output_path << "\n";
|
||||
|
||||
json root;
|
||||
if (cfg_.verbosity == Verbosity::xray) {
|
||||
root = build_xray();
|
||||
} else {
|
||||
root["movie"] = cfg_.movie_path;
|
||||
root["sample_fps"] = cfg_.sample_fps;
|
||||
root["anneal_sec"] = cfg_.anneal_sec;
|
||||
root["actors"] = build_epochs();
|
||||
if (cfg_.verbosity == Verbosity::standard)
|
||||
root["frames"] = build_standard();
|
||||
}
|
||||
|
||||
std::ofstream f(cfg_.output_path);
|
||||
if (!f.is_open()) {
|
||||
std::cerr << "[result_sink] ERROR: cannot write " << cfg_.output_path << "\n";
|
||||
return;
|
||||
}
|
||||
f << root.dump(2) << "\n";
|
||||
std::cerr << "[result_sink] done.\n";
|
||||
}
|
||||
|
||||
struct ActorWindow {
|
||||
std::string name, imdb_id;
|
||||
std::vector<std::pair<double, double>> scenes; // [start_sec, end_sec]
|
||||
};
|
||||
|
||||
// Core logic: merge per-frame detections into annealed [start, end] windows.
|
||||
std::vector<ActorWindow> build_actor_windows() {
|
||||
struct Info { std::string name, imdb_id; };
|
||||
std::map<int, Info> actor_info;
|
||||
std::map<int, std::vector<double>> timestamps;
|
||||
|
||||
for (const auto& frame : frames_) {
|
||||
for (const auto& ia : frame.visible_actors) {
|
||||
if (ia.actor_idx < 0) continue;
|
||||
actor_info[ia.actor_idx] = {ia.name, ia.imdb_id};
|
||||
timestamps[ia.actor_idx].push_back(frame.timestamp_sec);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<ActorWindow> result;
|
||||
for (auto& [idx, ts_vec] : timestamps) {
|
||||
ActorWindow aw;
|
||||
aw.name = actor_info[idx].name;
|
||||
aw.imdb_id = actor_info[idx].imdb_id;
|
||||
|
||||
double win_start = ts_vec[0], win_end = ts_vec[0];
|
||||
for (size_t i = 1; i < ts_vec.size(); ++i) {
|
||||
if (ts_vec[i] - win_end > cfg_.anneal_sec) {
|
||||
aw.scenes.push_back({win_start, win_end});
|
||||
win_start = ts_vec[i];
|
||||
}
|
||||
win_end = ts_vec[i];
|
||||
}
|
||||
aw.scenes.push_back({win_start, win_end});
|
||||
result.push_back(std::move(aw));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
json build_epochs() {
|
||||
json actors = json::array();
|
||||
for (const auto& aw : build_actor_windows()) {
|
||||
json windows = json::array();
|
||||
for (const auto& [s, e] : aw.scenes)
|
||||
windows.push_back({s, e});
|
||||
json ja;
|
||||
ja["name"] = aw.name;
|
||||
ja["imdb_id"] = aw.imdb_id;
|
||||
ja["scenes"] = std::move(windows);
|
||||
actors.push_back(std::move(ja));
|
||||
}
|
||||
return actors;
|
||||
}
|
||||
|
||||
// Jellyfin-Xray format: { "second": ["Actor", ...] }
|
||||
// Expands each annealed window into every integer second so coverage is dense
|
||||
// regardless of sample rate. Seconds between scenes have no key → overlay clears.
|
||||
json build_xray() {
|
||||
std::map<int, std::vector<std::string>> xray;
|
||||
for (const auto& aw : build_actor_windows()) {
|
||||
for (const auto& [start, end] : aw.scenes) {
|
||||
int t0 = static_cast<int>(std::floor(start));
|
||||
int t1 = static_cast<int>(std::ceil(end));
|
||||
for (int t = t0; t <= t1; ++t)
|
||||
xray[t].push_back(aw.name);
|
||||
}
|
||||
}
|
||||
|
||||
json root = json::object();
|
||||
for (const auto& [t, names] : xray)
|
||||
root[std::to_string(t)] = names;
|
||||
return root;
|
||||
}
|
||||
|
||||
json build_standard() {
|
||||
json frames = json::array();
|
||||
for (const auto& frame : frames_) {
|
||||
json jf;
|
||||
jf["t"] = frame.timestamp_sec;
|
||||
jf["identified"] = json::array();
|
||||
jf["unknowns"] = json::array();
|
||||
|
||||
for (const auto& ia : frame.visible_actors) {
|
||||
const auto& b = ia.bbox;
|
||||
json jbox = {b.x, b.y, b.width, b.height};
|
||||
|
||||
if (ia.actor_idx >= 0) {
|
||||
json ja;
|
||||
ja["name"] = ia.name;
|
||||
ja["imdb_id"] = ia.imdb_id;
|
||||
ja["similarity"] = ia.similarity;
|
||||
ja["track_id"] = ia.track_id;
|
||||
ja["bbox"] = jbox;
|
||||
jf["identified"].push_back(std::move(ja));
|
||||
} else {
|
||||
json ju;
|
||||
ju["bbox"] = jbox;
|
||||
ju["track_id"] = ia.track_id;
|
||||
ju["confidence"] = ia.similarity; // reuse field; 0 for unknowns
|
||||
jf["unknowns"].push_back(std::move(ju));
|
||||
}
|
||||
}
|
||||
frames.push_back(std::move(jf));
|
||||
}
|
||||
return frames;
|
||||
}
|
||||
|
||||
const Config& cfg_;
|
||||
std::atomic<bool>& done_;
|
||||
std::atomic<bool> written_{false};
|
||||
std::vector<SceneAnnotation> frames_;
|
||||
};
|
||||
Reference in New Issue
Block a user