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,115 @@
|
||||
#include "gallery_builder.hpp"
|
||||
#include "arcface_embedder.hpp"
|
||||
#include "face_utils.hpp"
|
||||
#include "ort_provider.hpp"
|
||||
#include "scrfd_decoder.hpp"
|
||||
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
// ── Parse "nm0000093_Brad_Pitt" → ("nm0000093", "Brad Pitt") ─────────────────
|
||||
static std::pair<std::string, std::string> parse_dir_name(const std::string& dirname) {
|
||||
auto pos = dirname.find('_');
|
||||
if (pos == std::string::npos) return {dirname, dirname};
|
||||
|
||||
std::string imdb_id = dirname.substr(0, pos);
|
||||
std::string raw = dirname.substr(pos + 1);
|
||||
std::string name;
|
||||
name.reserve(raw.size());
|
||||
for (char c : raw)
|
||||
name += (c == '_' ? ' ' : c);
|
||||
return {imdb_id, name};
|
||||
}
|
||||
|
||||
// ── Public API ────────────────────────────────────────────────────────────────
|
||||
|
||||
ActorGallery build_gallery(const BuildConfig& cfg) {
|
||||
const OrtProvider provider = detect_ort_provider();
|
||||
std::cerr << "[build_gallery] inference provider: " << provider_name(provider) << "\n";
|
||||
|
||||
SCRFDDecoder decoder(cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider);
|
||||
ArcFaceEmbedder arcface(cfg.arcface_model, provider);
|
||||
|
||||
ActorGallery gallery;
|
||||
|
||||
for (const auto& actor_dir : fs::directory_iterator(cfg.gallery_root)) {
|
||||
if (!actor_dir.is_directory()) continue;
|
||||
|
||||
auto [imdb_id, name] = parse_dir_name(actor_dir.path().filename().string());
|
||||
std::cerr << "[build_gallery] " << name << " (" << imdb_id << ")\n";
|
||||
|
||||
ActorGallery::Actor actor;
|
||||
actor.imdb_id = imdb_id;
|
||||
actor.name = name;
|
||||
|
||||
static const std::vector<std::string> kExts{".jpg", ".jpeg", ".png", ".webp"};
|
||||
for (const auto& img_file : fs::directory_iterator(actor_dir.path())) {
|
||||
if (!img_file.is_regular_file()) continue;
|
||||
std::string ext = img_file.path().extension().string();
|
||||
std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
|
||||
if (std::find(kExts.begin(), kExts.end(), ext) == kExts.end()) continue;
|
||||
|
||||
cv::Mat img = cv::imread(img_file.path().string());
|
||||
if (img.empty()) {
|
||||
std::cerr << " [skip] cannot read " << img_file.path().filename() << "\n";
|
||||
continue;
|
||||
}
|
||||
|
||||
if (cfg.max_side > 0) {
|
||||
const int big = std::max(img.cols, img.rows);
|
||||
if (big > cfg.max_side) {
|
||||
const double s = static_cast<double>(cfg.max_side) / big;
|
||||
cv::resize(img, img, {}, s, s, cv::INTER_AREA);
|
||||
}
|
||||
}
|
||||
|
||||
auto faces = decoder.detect(img);
|
||||
|
||||
if (faces.empty()) {
|
||||
std::cerr << " [skip] no face: " << img_file.path().filename() << "\n";
|
||||
continue;
|
||||
}
|
||||
if (faces.size() > 1) {
|
||||
std::cerr << " [warn] " << faces.size() << " faces, using highest confidence: "
|
||||
<< img_file.path().filename() << "\n";
|
||||
}
|
||||
|
||||
const auto& best = *std::max_element(
|
||||
faces.begin(), faces.end(),
|
||||
[](const DetectedFace& a, const DetectedFace& b) {
|
||||
return a.confidence < b.confidence;
|
||||
});
|
||||
|
||||
cv::Mat crop = align_face(img, best.landmarks);
|
||||
if (crop.empty()) {
|
||||
std::cerr << " [skip] alignment failed: " << img_file.path().filename() << "\n";
|
||||
continue;
|
||||
}
|
||||
|
||||
Embedding emb = arcface.embed_one(crop);
|
||||
actor.embeddings.push_back(emb);
|
||||
actor.source_images.push_back(img_file.path().filename().string());
|
||||
|
||||
std::cerr << " [ok] " << img_file.path().filename()
|
||||
<< " conf=" << best.confidence << "\n";
|
||||
}
|
||||
|
||||
if (actor.embeddings.empty()) {
|
||||
std::cerr << " [warn] no valid embeddings for " << name << " — skipped\n";
|
||||
continue;
|
||||
}
|
||||
|
||||
std::cerr << " → " << actor.embeddings.size() << " embeddings\n";
|
||||
gallery.actors.push_back(std::move(actor));
|
||||
}
|
||||
|
||||
std::cerr << "[build_gallery] total: " << gallery.actors.size() << " actors\n";
|
||||
return gallery;
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include <string>
|
||||
|
||||
// Build an ActorGallery from a directory tree:
|
||||
//
|
||||
// gallery_root/
|
||||
// nm0000093_Brad_Pitt/
|
||||
// img1.jpg
|
||||
// img2.jpg
|
||||
// ...
|
||||
// nm0000129_Cate_Blanchett/
|
||||
// ...
|
||||
//
|
||||
// Each subdirectory name is parsed as "<imdb_id>_<Name_With_Underscores>".
|
||||
// For every image:
|
||||
// 1. Detect face with SCRFD-500MF (expect exactly one; warn and skip if 0 or >1).
|
||||
// 2. Align with ArcFace 5-point transform → 112×112 crop.
|
||||
// 3. Embed with ArcFace ONNX → 512-dim L2-normalised embedding.
|
||||
// All embeddings are stored (best-of-N match at query time).
|
||||
//
|
||||
// Returns a gallery ready to pass to save_gallery() / IdentityMatcherFunc.
|
||||
|
||||
struct BuildConfig {
|
||||
std::string gallery_root; // directory tree described above
|
||||
std::string detector_model;
|
||||
std::string arcface_model;
|
||||
float detector_conf{0.5f};
|
||||
float detector_nms{0.4f};
|
||||
int max_side{500}; // downscale source images to this max dimension
|
||||
// before detection — TMDB portraits are ~2k px,
|
||||
// SCRFD trains on smaller faces and detection
|
||||
// confidence drops on huge inputs. 0 = disabled.
|
||||
};
|
||||
|
||||
ActorGallery build_gallery(const BuildConfig& cfg);
|
||||
@@ -0,0 +1,118 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
// ── GalleryCalibration ────────────────────────────────────────────────────────
|
||||
// Platt-style sigmoid calibration: P(match) = σ(a · similarity + b)
|
||||
// where similarity = cosine similarity ∈ [-1, 1] (dot product of L2-normalised
|
||||
// ArcFace embeddings).
|
||||
//
|
||||
// Fitted from intra-class (positive) and inter-class (negative) pairs built
|
||||
// from the gallery reference embeddings. When calibration is invalid (too few
|
||||
// positive pairs), the caller should fall back to the raw threshold.
|
||||
struct GalleryCalibration {
|
||||
float a{10.f}; // scale (positive → higher similarity → higher probability)
|
||||
float b{-5.f}; // bias (decision boundary at similarity = -b/a)
|
||||
bool valid{false};
|
||||
|
||||
// P(match | sim) using the balanced-prior calibration.
|
||||
// Pass log_prior_odds = log(p0/(1-p0)) to adjust for a real base-rate prior p0:
|
||||
// P(match | sim, p0) = σ(a·sim + b + log(p0/(1-p0)))
|
||||
float probability(float similarity, float log_prior_odds = 0.f) const {
|
||||
float z = a * similarity + b + log_prior_odds;
|
||||
if (z >= 0.f) return 1.f / (1.f + std::exp(-z));
|
||||
float e = std::exp(z);
|
||||
return e / (1.f + e);
|
||||
}
|
||||
|
||||
// Similarity at which P(match, prior) == p
|
||||
float boundary_at(float p = 0.5f, float log_prior_odds = 0.f) const {
|
||||
return (std::log(p / (1.f - p)) - b - log_prior_odds) / a;
|
||||
}
|
||||
};
|
||||
|
||||
// Fit a logistic sigmoid to gallery pair similarities.
|
||||
// Positive pairs: same actor, different reference images.
|
||||
// Negative pairs: different actors (all cross-actor embedding pairs).
|
||||
// Class weights balance the (typically skewed) pos/neg ratio.
|
||||
// Requires ≥2 positive pairs and ≥1 negative pair.
|
||||
inline GalleryCalibration calibrate_gallery(
|
||||
const std::vector<Embedding>& flat_emb,
|
||||
const std::vector<int>& flat_actor)
|
||||
{
|
||||
const int n = static_cast<int>(flat_emb.size());
|
||||
|
||||
std::vector<float> X; // cosine similarities
|
||||
std::vector<float> Y; // labels: 1 = same actor, 0 = different
|
||||
|
||||
for (int i = 0; i < n; ++i) {
|
||||
for (int j = i + 1; j < n; ++j) {
|
||||
float dot = 0.f;
|
||||
for (int k = 0; k < 512; ++k)
|
||||
dot += flat_emb[i][k] * flat_emb[j][k];
|
||||
X.push_back(dot);
|
||||
Y.push_back(flat_actor[i] == flat_actor[j] ? 1.f : 0.f);
|
||||
}
|
||||
}
|
||||
|
||||
int n_pos = 0;
|
||||
for (float y : Y) if (y > 0.5f) ++n_pos;
|
||||
int n_neg = static_cast<int>(Y.size()) - n_pos;
|
||||
|
||||
if (n_pos < 2 || n_neg < 1) {
|
||||
std::cerr << "[calibration] insufficient pairs (+" << n_pos
|
||||
<< "/-" << n_neg << ") — calibration skipped\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
// Class weights to handle pos/neg imbalance
|
||||
float total = static_cast<float>(Y.size());
|
||||
float w_pos = total / (2.f * n_pos);
|
||||
float w_neg = total / (2.f * n_neg);
|
||||
|
||||
// Gradient descent logistic regression (2 parameters: a, b)
|
||||
float a = 10.f, b = -5.f;
|
||||
constexpr float lr = 0.05f;
|
||||
constexpr int max_iter = 20000;
|
||||
constexpr float tol = 1e-7f;
|
||||
|
||||
for (int iter = 0; iter < max_iter; ++iter) {
|
||||
float da = 0.f, db = 0.f;
|
||||
for (int i = 0; i < static_cast<int>(X.size()); ++i) {
|
||||
float z = a * X[i] + b;
|
||||
float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z))
|
||||
: std::exp(z) / (1.f + std::exp(z));
|
||||
float err = sig - Y[i];
|
||||
float w = (Y[i] > 0.5f) ? w_pos : w_neg;
|
||||
da += w * err * X[i];
|
||||
db += w * err;
|
||||
}
|
||||
da /= total;
|
||||
db /= total;
|
||||
a -= lr * da;
|
||||
b -= lr * db;
|
||||
if (da * da + db * db < tol * tol) break;
|
||||
}
|
||||
|
||||
// Training accuracy at P=0.5 decision boundary
|
||||
int correct = 0;
|
||||
for (int i = 0; i < static_cast<int>(X.size()); ++i) {
|
||||
float z = a * X[i] + b;
|
||||
float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z))
|
||||
: std::exp(z) / (1.f + std::exp(z));
|
||||
if ((sig > 0.5f) == (Y[i] > 0.5f)) ++correct;
|
||||
}
|
||||
float acc = 100.f * correct / static_cast<float>(Y.size());
|
||||
|
||||
GalleryCalibration cal{a, b, true};
|
||||
std::cerr << "[calibration] sigmoid fitted:"
|
||||
<< " a=" << a << " b=" << b
|
||||
<< " boundary(P=0.5)=sim" << cal.boundary_at(0.5f)
|
||||
<< " pairs=" << Y.size()
|
||||
<< " (+" << n_pos << "/-" << n_neg << ")"
|
||||
<< " train_acc=" << acc << "%\n";
|
||||
return cal;
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
#include "gallery_store.hpp"
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <fstream>
|
||||
#include <stdexcept>
|
||||
|
||||
using json = nlohmann::json;
|
||||
|
||||
ActorGallery load_gallery(const std::string& path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.is_open())
|
||||
throw std::runtime_error("load_gallery: cannot open " + path);
|
||||
|
||||
json j;
|
||||
f >> j;
|
||||
|
||||
ActorGallery gallery;
|
||||
for (const auto& ja : j.at("actors")) {
|
||||
ActorGallery::Actor actor;
|
||||
actor.imdb_id = ja.at("imdb_id").get<std::string>();
|
||||
actor.name = ja.at("name").get<std::string>();
|
||||
|
||||
if (ja.contains("source_images"))
|
||||
actor.source_images = ja.at("source_images").get<std::vector<std::string>>();
|
||||
|
||||
for (const auto& je : ja.at("embeddings")) {
|
||||
Embedding emb = je.get<Embedding>();
|
||||
actor.embeddings.push_back(emb);
|
||||
}
|
||||
|
||||
gallery.actors.push_back(std::move(actor));
|
||||
}
|
||||
return gallery;
|
||||
}
|
||||
|
||||
void save_gallery(const std::string& path, const ActorGallery& gallery) {
|
||||
json j;
|
||||
j["actors"] = json::array();
|
||||
|
||||
for (const auto& actor : gallery.actors) {
|
||||
json ja;
|
||||
ja["imdb_id"] = actor.imdb_id;
|
||||
ja["name"] = actor.name;
|
||||
ja["source_images"] = actor.source_images;
|
||||
|
||||
ja["embeddings"] = json::array();
|
||||
for (const auto& emb : actor.embeddings) {
|
||||
ja["embeddings"].push_back(
|
||||
std::vector<float>(emb.begin(), emb.end()));
|
||||
}
|
||||
j["actors"].push_back(std::move(ja));
|
||||
}
|
||||
|
||||
std::ofstream f(path);
|
||||
if (!f.is_open())
|
||||
throw std::runtime_error("save_gallery: cannot write " + path);
|
||||
f << j.dump(2) << "\n";
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
#include <string>
|
||||
|
||||
// Load/save the actor gallery from/to a JSON file.
|
||||
//
|
||||
// JSON format:
|
||||
// {
|
||||
// "actors": [
|
||||
// {
|
||||
// "imdb_id": "nm0000093",
|
||||
// "name": "Brad Pitt",
|
||||
// "source_images": ["img1.jpg", "img2.jpg"],
|
||||
// "embeddings": [[0.012, -0.034, ...], ...] // one 512-float array per image
|
||||
// }
|
||||
// ]
|
||||
// }
|
||||
|
||||
ActorGallery load_gallery(const std::string& path);
|
||||
void save_gallery(const std::string& path, const ActorGallery& gallery);
|
||||
Reference in New Issue
Block a user