faster calibration curve generation
jellyfin intergration
This commit is contained in:
@@ -0,0 +1,104 @@
|
||||
#pragma once
|
||||
// FaceEmbedderEngine — load SCRFD + ArcFace once, embed many images.
|
||||
//
|
||||
// Extracted from embed_faces.cpp so the same detect→align→embed pipeline can
|
||||
// be driven from a long-lived process (the sae_embed Python module) instead
|
||||
// of a fresh CLI invocation per image, which would reload both ONNX sessions
|
||||
// every time.
|
||||
|
||||
#include "arcface_embedder.hpp"
|
||||
#include "face_utils.hpp"
|
||||
#include "ort_provider.hpp"
|
||||
#include "scrfd_decoder.hpp"
|
||||
#include "types.hpp"
|
||||
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
struct FaceEmbedResult {
|
||||
bool ok{false};
|
||||
std::string error;
|
||||
Embedding embedding{};
|
||||
float confidence{0.f};
|
||||
float bbox[4]{}; // x, y, w, h
|
||||
std::array<cv::Point2f, 5> landmarks{};
|
||||
};
|
||||
|
||||
class FaceEmbedderEngine {
|
||||
public:
|
||||
FaceEmbedderEngine(const std::string& detector_model,
|
||||
const std::string& arcface_model,
|
||||
float conf = 0.5f, float nms = 0.4f, int max_side = 500)
|
||||
: max_side_(max_side)
|
||||
{
|
||||
const OrtProvider provider = detect_ort_provider();
|
||||
std::cerr << "[FaceEmbedderEngine] inference provider: "
|
||||
<< provider_name(provider) << "\n";
|
||||
detector_ = std::make_unique<SCRFDDecoder>(detector_model, conf, nms, provider);
|
||||
embedder_ = std::make_unique<ArcFaceEmbedder>(arcface_model, provider);
|
||||
}
|
||||
|
||||
FaceEmbedResult embed_path(const std::string& path) const {
|
||||
cv::Mat img = cv::imread(path);
|
||||
if (img.empty()) {
|
||||
FaceEmbedResult res;
|
||||
res.error = "cannot read image";
|
||||
return res;
|
||||
}
|
||||
return embed_mat(img);
|
||||
}
|
||||
|
||||
FaceEmbedResult embed_mat(cv::Mat img) const {
|
||||
FaceEmbedResult res;
|
||||
|
||||
if (max_side_ > 0) {
|
||||
const int big = std::max(img.cols, img.rows);
|
||||
if (big > max_side_) {
|
||||
const double s = static_cast<double>(max_side_) / big;
|
||||
cv::resize(img, img, {}, s, s, cv::INTER_AREA);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<DetectedFace> faces = detector_->detect(img);
|
||||
if (faces.empty()) {
|
||||
res.error = "no face detected";
|
||||
return res;
|
||||
}
|
||||
if (faces.size() > 1)
|
||||
std::cerr << "[warn] " << faces.size()
|
||||
<< " faces detected, using highest-confidence one\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()) {
|
||||
res.error = "alignment failed";
|
||||
return res;
|
||||
}
|
||||
|
||||
res.ok = true;
|
||||
res.embedding = embedder_->embed_one(crop);
|
||||
res.confidence = best.confidence;
|
||||
res.bbox[0] = best.bbox.x;
|
||||
res.bbox[1] = best.bbox.y;
|
||||
res.bbox[2] = best.bbox.width;
|
||||
res.bbox[3] = best.bbox.height;
|
||||
res.landmarks = best.landmarks;
|
||||
return res;
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<SCRFDDecoder> detector_;
|
||||
std::unique_ptr<ArcFaceEmbedder> embedder_;
|
||||
int max_side_;
|
||||
};
|
||||
@@ -1,10 +1,24 @@
|
||||
#pragma once
|
||||
#include "types.hpp"
|
||||
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <opencv2/core/ocl.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
|
||||
// ── GalleryCalibration ────────────────────────────────────────────────────────
|
||||
// Platt-style sigmoid calibration: P(match) = σ(a · similarity + b)
|
||||
// where similarity = cosine similarity ∈ [-1, 1] (dot product of L2-normalised
|
||||
@@ -43,24 +57,147 @@ 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());
|
||||
constexpr int kMinEmbeddingsForPositive = 5;
|
||||
constexpr float kDedupSimThreshold = 1.f - 1e-7f; // sim above this = duplicate
|
||||
constexpr int kHistBins = 200;
|
||||
|
||||
std::vector<float> X; // cosine similarities
|
||||
std::vector<float> Y; // labels: 1 = same actor, 0 = different
|
||||
// ── Per-actor dedup, then eligibility filter ────────────────────────────
|
||||
// Drop near-identical duplicate embeddings within each actor (e.g. the
|
||||
// same source image embedded twice). Actors left with fewer than
|
||||
// kMinEmbeddingsForPositive distinct embeddings can't supply meaningful
|
||||
// same-actor (positive) pairs, but still contribute negative
|
||||
// (cross-actor) pairs for the p(unknown) side of the fit.
|
||||
int n_actors = 0;
|
||||
for (int a : flat_actor) n_actors = std::max(n_actors, a + 1);
|
||||
|
||||
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);
|
||||
std::vector<std::vector<Embedding>> by_actor(n_actors);
|
||||
for (int i = 0; i < static_cast<int>(flat_emb.size()); ++i)
|
||||
by_actor[flat_actor[i]].push_back(flat_emb[i]);
|
||||
|
||||
std::vector<Embedding> flat_emb_dedup;
|
||||
std::vector<int> flat_actor_dedup;
|
||||
std::vector<bool> actor_eligible(n_actors, false);
|
||||
int n_eligible = 0;
|
||||
|
||||
for (int ai = 0; ai < n_actors; ++ai) {
|
||||
std::vector<Embedding> kept;
|
||||
for (const auto& e : by_actor[ai]) {
|
||||
bool dup = false;
|
||||
for (const auto& k : kept) {
|
||||
if (cosine_similarity(e, k) > kDedupSimThreshold) { dup = true; break; }
|
||||
}
|
||||
if (!dup) kept.push_back(e);
|
||||
}
|
||||
if (static_cast<int>(kept.size()) >= kMinEmbeddingsForPositive) {
|
||||
actor_eligible[ai] = true;
|
||||
++n_eligible;
|
||||
}
|
||||
for (auto& e : kept) {
|
||||
flat_emb_dedup.push_back(e);
|
||||
flat_actor_dedup.push_back(ai);
|
||||
}
|
||||
}
|
||||
|
||||
int n_pos = 0;
|
||||
for (float y : Y) if (y > 0.5f) ++n_pos;
|
||||
int n_neg = static_cast<int>(Y.size()) - n_pos;
|
||||
const int n = static_cast<int>(flat_emb_dedup.size());
|
||||
std::cerr << "[calibration] dedup: " << flat_emb.size() << " -> " << n
|
||||
<< " embeddings (" << n_eligible << "/" << n_actors
|
||||
<< " actors have >= " << kMinEmbeddingsForPositive
|
||||
<< " distinct embeddings, eligible for positive pairs)\n";
|
||||
|
||||
// Pairwise cosine-similarity matrix S = E * E^T, computed via cv::gemm
|
||||
// (BLAS/SIMD on CPU) instead of a naive O(n^2 * 512) scalar loop.
|
||||
//
|
||||
// SAE_CALIB_GEMM controls the backend:
|
||||
// "cpu" - cv::Mat gemm only (default if OpenCL unavailable)
|
||||
// "gpu" - cv::UMat gemm via OpenCL, falls back to CPU if unavailable
|
||||
// "bench" - run both and report timings + max abs diff (default)
|
||||
cv::Mat E(n, 512, CV_32F);
|
||||
for (int i = 0; i < n; ++i)
|
||||
std::memcpy(E.ptr<float>(i), flat_emb_dedup[i].data(), 512 * sizeof(float));
|
||||
|
||||
std::string mode = [] {
|
||||
const char* env = std::getenv("SAE_CALIB_GEMM");
|
||||
return env ? std::string(env) : std::string("bench");
|
||||
}();
|
||||
|
||||
bool have_ocl = cv::ocl::haveOpenCL();
|
||||
std::cerr << "[calibration] OpenCL available: " << (have_ocl ? "yes" : "no");
|
||||
if (have_ocl) {
|
||||
auto dev = cv::ocl::Device::getDefault();
|
||||
std::cerr << " device=\"" << dev.name() << "\""
|
||||
<< " type=" << (dev.type() == cv::ocl::Device::TYPE_GPU ? "GPU"
|
||||
: dev.type() == cv::ocl::Device::TYPE_CPU ? "CPU"
|
||||
: "OTHER");
|
||||
}
|
||||
std::cerr << " (SAE_CALIB_GEMM=" << mode << ")\n";
|
||||
|
||||
std::cerr << "[calibration] similarity matrix: " << n << "x512 ("
|
||||
<< (n * 512LL * sizeof(float)) / (1024 * 1024) << " MB input, "
|
||||
<< (n * (long long)n * sizeof(float)) / (1024 * 1024)
|
||||
<< " MB output)\n" << std::flush;
|
||||
|
||||
cv::Mat S;
|
||||
bool have_cpu = false, have_gpu = false;
|
||||
cv::Mat S_cpu, S_gpu_host;
|
||||
|
||||
if (mode == "cpu" || mode == "bench") {
|
||||
std::cerr << "[calibration] running cv::gemm on CPU...\n" << std::flush;
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
cv::gemm(E, E, 1.0, cv::noArray(), 0.0, S_cpu, cv::GEMM_2_T);
|
||||
auto t1 = std::chrono::steady_clock::now();
|
||||
double secs = std::chrono::duration<double>(t1 - t0).count();
|
||||
std::cerr << "[calibration] cv::gemm CPU: " << secs << "s for "
|
||||
<< n << "x" << n << " similarity matrix\n";
|
||||
have_cpu = true;
|
||||
}
|
||||
|
||||
if ((mode == "gpu" || mode == "bench") && have_ocl) {
|
||||
std::cerr << "[calibration] running cv::gemm on OpenCL device...\n" << std::flush;
|
||||
cv::UMat E_gpu, S_gpu;
|
||||
E.copyTo(E_gpu);
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
cv::gemm(E_gpu, E_gpu, 1.0, cv::noArray(), 0.0, S_gpu, cv::GEMM_2_T);
|
||||
S_gpu_host = S_gpu.getMat(cv::ACCESS_READ).clone();
|
||||
auto t1 = std::chrono::steady_clock::now();
|
||||
double secs = std::chrono::duration<double>(t1 - t0).count();
|
||||
std::cerr << "[calibration] cv::gemm OpenCL: " << secs << "s for "
|
||||
<< n << "x" << n << " similarity matrix";
|
||||
if (have_cpu) std::cerr << " max|diff|=" << cv::norm(S_cpu, S_gpu_host, cv::NORM_INF);
|
||||
std::cerr << "\n";
|
||||
have_gpu = true;
|
||||
} else if (mode == "gpu") {
|
||||
std::cerr << "[calibration] SAE_CALIB_GEMM=gpu requested but OpenCL is unavailable\n";
|
||||
}
|
||||
|
||||
if (mode == "gpu" && have_gpu) S = S_gpu_host;
|
||||
else if (have_cpu) S = S_cpu;
|
||||
else if (have_gpu) S = S_gpu_host;
|
||||
else throw std::runtime_error("calibrate_gallery: no gemm backend produced a result");
|
||||
|
||||
// ── Histogram of pairwise similarities ──────────────────────────────────
|
||||
// Instead of storing one (similarity, label) sample per pair (~n^2/2
|
||||
// entries — too many for the gradient descent below), bucket pairs into
|
||||
// kHistBins bins over sim ∈ [-1, 1] and fit the sigmoid against the
|
||||
// per-bin (positive_count, negative_count).
|
||||
std::vector<double> pos_hist(kHistBins, 0.0), neg_hist(kHistBins, 0.0);
|
||||
constexpr float bin_width = 2.f / kHistBins;
|
||||
|
||||
for (int i = 0; i < n; ++i) {
|
||||
const float* row = S.ptr<float>(i);
|
||||
for (int j = i + 1; j < n; ++j) {
|
||||
int bin = static_cast<int>((row[j] + 1.f) / bin_width);
|
||||
bin = std::clamp(bin, 0, kHistBins - 1);
|
||||
if (flat_actor_dedup[i] == flat_actor_dedup[j]) {
|
||||
if (actor_eligible[flat_actor_dedup[i]]) pos_hist[bin] += 1.0;
|
||||
// same actor but ineligible (< kMinEmbeddingsForPositive) — skip
|
||||
} else {
|
||||
neg_hist[bin] += 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
double n_pos = 0.0, n_neg = 0.0;
|
||||
for (int b = 0; b < kHistBins; ++b) { n_pos += pos_hist[b]; n_neg += neg_hist[b]; }
|
||||
|
||||
if (n_pos < 2 || n_neg < 1) {
|
||||
std::cerr << "[calibration] insufficient pairs (+" << n_pos
|
||||
@@ -69,50 +206,192 @@ inline GalleryCalibration calibrate_gallery(
|
||||
}
|
||||
|
||||
// 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);
|
||||
double total = n_pos + n_neg;
|
||||
double w_pos = total / (2.0 * n_pos);
|
||||
double w_neg = total / (2.0 * n_neg);
|
||||
|
||||
// Gradient descent logistic regression (2 parameters: a, b)
|
||||
float a = 10.f, b = -5.f;
|
||||
std::vector<float> bin_center(kHistBins);
|
||||
for (int b = 0; b < kHistBins; ++b) bin_center[b] = -1.f + (b + 0.5f) * bin_width;
|
||||
|
||||
// Gradient descent logistic regression (2 parameters: a, bias)
|
||||
float a = 10.f, bias = -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;
|
||||
double da = 0.0, db = 0.0;
|
||||
for (int b = 0; b < kHistBins; ++b) {
|
||||
float x = bin_center[b];
|
||||
float z = a * x + bias;
|
||||
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;
|
||||
if (pos_hist[b] > 0.0) {
|
||||
double err = (sig - 1.f) * w_pos * pos_hist[b];
|
||||
da += err * x;
|
||||
db += err;
|
||||
}
|
||||
if (neg_hist[b] > 0.0) {
|
||||
double err = sig * w_neg * neg_hist[b];
|
||||
da += err * x;
|
||||
db += err;
|
||||
}
|
||||
}
|
||||
da /= total;
|
||||
db /= total;
|
||||
a -= lr * da;
|
||||
b -= lr * db;
|
||||
a -= lr * static_cast<float>(da);
|
||||
bias -= lr * static_cast<float>(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;
|
||||
double correct = 0.0;
|
||||
for (int b = 0; b < kHistBins; ++b) {
|
||||
float z = a * bin_center[b] + bias;
|
||||
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;
|
||||
correct += (sig > 0.5f) ? pos_hist[b] : neg_hist[b];
|
||||
}
|
||||
float acc = 100.f * correct / static_cast<float>(Y.size());
|
||||
double acc = 100.0 * correct / total;
|
||||
|
||||
GalleryCalibration cal{a, b, true};
|
||||
GalleryCalibration cal{a, bias, true};
|
||||
std::cerr << "[calibration] sigmoid fitted:"
|
||||
<< " a=" << a << " b=" << b
|
||||
<< " a=" << a << " b=" << bias
|
||||
<< " boundary(P=0.5)=sim" << cal.boundary_at(0.5f)
|
||||
<< " pairs=" << Y.size()
|
||||
<< " (+" << n_pos << "/-" << n_neg << ")"
|
||||
<< " pairs=" << static_cast<long long>(total)
|
||||
<< " (+" << static_cast<long long>(n_pos) << "/-" << static_cast<long long>(n_neg) << ")"
|
||||
<< " bins=" << kHistBins
|
||||
<< " train_acc=" << acc << "%\n";
|
||||
return cal;
|
||||
}
|
||||
|
||||
// Dumps the fitted P(match | similarity) sigmoid as both a CSV table
|
||||
// (similarity, p_match) and a PNG plot, to <base_path>.csv / <base_path>.png.
|
||||
inline void save_calibration_curve(const GalleryCalibration& cal, const std::string& base_path) {
|
||||
if (!cal.valid) return;
|
||||
|
||||
constexpr int kSamples = 200;
|
||||
|
||||
std::ofstream csv(base_path + ".csv");
|
||||
if (csv.is_open()) {
|
||||
csv << "similarity,p_match\n";
|
||||
for (int i = 0; i <= kSamples; ++i) {
|
||||
float sim = -1.f + i * (2.f / kSamples);
|
||||
csv << sim << "," << cal.probability(sim) << "\n";
|
||||
}
|
||||
}
|
||||
|
||||
constexpr int W = 800, H = 600, M = 50;
|
||||
cv::Mat img(H, W, CV_8UC3, cv::Scalar(255, 255, 255));
|
||||
cv::line(img, {M, H - M}, {W - M, H - M}, {0, 0, 0}, 1); // x-axis: similarity [-1,1]
|
||||
cv::line(img, {M, M}, {M, H - M}, {0, 0, 0}, 1); // y-axis: P(match) [0,1]
|
||||
|
||||
auto to_point = [&](float sim, float p) {
|
||||
int x = M + static_cast<int>((sim + 1.f) * 0.5f * (W - 2 * M));
|
||||
int y = (H - M) - static_cast<int>(p * (H - 2 * M));
|
||||
return cv::Point(x, y);
|
||||
};
|
||||
|
||||
// P=0.5 reference line and decision boundary
|
||||
float boundary = cal.boundary_at(0.5f);
|
||||
cv::line(img, to_point(-1.f, 0.5f), to_point(1.f, 0.5f), {200, 200, 200}, 1);
|
||||
if (boundary >= -1.f && boundary <= 1.f)
|
||||
cv::line(img, to_point(boundary, 0.f), to_point(boundary, 1.f), {200, 200, 200}, 1);
|
||||
|
||||
cv::Point prev = to_point(-1.f, cal.probability(-1.f));
|
||||
for (int i = 1; i <= kSamples; ++i) {
|
||||
float sim = -1.f + i * (2.f / kSamples);
|
||||
cv::Point pt = to_point(sim, cal.probability(sim));
|
||||
cv::line(img, prev, pt, {255, 0, 0}, 2);
|
||||
prev = pt;
|
||||
}
|
||||
|
||||
cv::imwrite(base_path + ".png", img);
|
||||
std::cerr << "[calibration] saved curve to " << base_path << ".csv / " << base_path << ".png\n";
|
||||
}
|
||||
|
||||
// FNV-1a 64-bit hash over the gallery's reference embeddings and actor
|
||||
// assignments, used to detect whether a cached calibration is still valid.
|
||||
inline uint64_t hash_gallery_embeddings(
|
||||
const std::vector<Embedding>& flat_emb,
|
||||
const std::vector<int>& flat_actor)
|
||||
{
|
||||
uint64_t h = 1469598103934665603ULL;
|
||||
constexpr uint64_t prime = 1099511628211ULL;
|
||||
auto mix = [&](const void* data, size_t n) {
|
||||
const auto* p = static_cast<const unsigned char*>(data);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
h ^= p[i];
|
||||
h *= prime;
|
||||
}
|
||||
};
|
||||
|
||||
uint64_t n = flat_emb.size();
|
||||
mix(&n, sizeof(n));
|
||||
for (const auto& emb : flat_emb)
|
||||
mix(emb.data(), emb.size() * sizeof(float));
|
||||
for (int a : flat_actor)
|
||||
mix(&a, sizeof(a));
|
||||
return h;
|
||||
}
|
||||
|
||||
// Calibrates the gallery, caching the fitted (a, b, valid) result on disk
|
||||
// keyed by a hash of the reference embeddings. The O(n^2) pairwise fit only
|
||||
// re-runs when the gallery's embeddings/actor assignments actually change.
|
||||
inline GalleryCalibration calibrate_gallery_cached(
|
||||
const std::vector<Embedding>& flat_emb,
|
||||
const std::vector<int>& flat_actor,
|
||||
const std::string& cache_path)
|
||||
{
|
||||
uint64_t hash = hash_gallery_embeddings(flat_emb, flat_actor);
|
||||
|
||||
std::string base_path = cache_path;
|
||||
constexpr std::string_view kJsonExt = ".json";
|
||||
if (base_path.size() >= kJsonExt.size() &&
|
||||
base_path.compare(base_path.size() - kJsonExt.size(), kJsonExt.size(), kJsonExt) == 0)
|
||||
base_path.resize(base_path.size() - kJsonExt.size());
|
||||
|
||||
std::ifstream in(cache_path);
|
||||
if (in.is_open()) {
|
||||
try {
|
||||
nlohmann::json j;
|
||||
in >> j;
|
||||
if (j.at("hash").get<uint64_t>() == hash) {
|
||||
GalleryCalibration cal;
|
||||
cal.a = j.at("a").get<float>();
|
||||
cal.b = j.at("b").get<float>();
|
||||
cal.valid = j.at("valid").get<bool>();
|
||||
std::cerr << "[calibration] using cached calibration from "
|
||||
<< cache_path << " (a=" << cal.a << " b=" << cal.b
|
||||
<< " valid=" << cal.valid << ")\n";
|
||||
save_calibration_curve(cal, base_path);
|
||||
return cal;
|
||||
}
|
||||
std::cerr << "[calibration] cache at " << cache_path
|
||||
<< " is stale, recomputing\n";
|
||||
} catch (const std::exception&) {
|
||||
std::cerr << "[calibration] cache at " << cache_path
|
||||
<< " is unreadable, recomputing\n";
|
||||
}
|
||||
}
|
||||
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
GalleryCalibration cal = calibrate_gallery(flat_emb, flat_actor);
|
||||
auto t1 = std::chrono::steady_clock::now();
|
||||
double secs = std::chrono::duration<double>(t1 - t0).count();
|
||||
std::cerr << "[calibration] fit took " << secs << "s for "
|
||||
<< flat_emb.size() << " embeddings\n";
|
||||
|
||||
nlohmann::json j;
|
||||
j["hash"] = hash;
|
||||
j["a"] = cal.a;
|
||||
j["b"] = cal.b;
|
||||
j["valid"] = cal.valid;
|
||||
j["fit_secs"] = secs;
|
||||
std::ofstream out(cache_path);
|
||||
if (out.is_open()) out << j.dump(2) << "\n";
|
||||
|
||||
save_calibration_curve(cal, base_path);
|
||||
|
||||
return cal;
|
||||
}
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
#include "gallery_store.hpp"
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <chrono>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
|
||||
using json = nlohmann::json;
|
||||
@@ -11,14 +13,24 @@ ActorGallery load_gallery(const std::string& path) {
|
||||
if (!f.is_open())
|
||||
throw std::runtime_error("load_gallery: cannot open " + path);
|
||||
|
||||
std::cerr << "[gallery] loading " << path << "..." << std::flush;
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
|
||||
json j;
|
||||
f >> j;
|
||||
|
||||
auto t1 = std::chrono::steady_clock::now();
|
||||
std::cerr << " parsed JSON in "
|
||||
<< std::chrono::duration<double>(t1 - t0).count() << "s\n";
|
||||
|
||||
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>();
|
||||
actor.imdb_id = ja.value("imdb_id", "");
|
||||
actor.tmdb_id = ja.value("tmdb_id", "");
|
||||
// older make_jellyfin_gallery.py galleries used "jellyfin_person_id"
|
||||
actor.jellyfin_id = ja.value("jellyfin_id", ja.value("jellyfin_person_id", ""));
|
||||
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>>();
|
||||
@@ -30,6 +42,14 @@ ActorGallery load_gallery(const std::string& path) {
|
||||
|
||||
gallery.actors.push_back(std::move(actor));
|
||||
}
|
||||
|
||||
size_t n_emb = 0;
|
||||
for (const auto& actor : gallery.actors) n_emb += actor.embeddings.size();
|
||||
auto t2 = std::chrono::steady_clock::now();
|
||||
std::cerr << "[gallery] built " << gallery.actors.size() << " actors / "
|
||||
<< n_emb << " embeddings in "
|
||||
<< std::chrono::duration<double>(t2 - t1).count() << "s\n";
|
||||
|
||||
return gallery;
|
||||
}
|
||||
|
||||
@@ -40,6 +60,8 @@ void save_gallery(const std::string& path, const ActorGallery& gallery) {
|
||||
for (const auto& actor : gallery.actors) {
|
||||
json ja;
|
||||
ja["imdb_id"] = actor.imdb_id;
|
||||
ja["tmdb_id"] = actor.tmdb_id;
|
||||
ja["jellyfin_id"] = actor.jellyfin_id;
|
||||
ja["name"] = actor.name;
|
||||
ja["source_images"] = actor.source_images;
|
||||
|
||||
|
||||
@@ -8,7 +8,9 @@
|
||||
// {
|
||||
// "actors": [
|
||||
// {
|
||||
// "imdb_id": "nm0000093",
|
||||
// "imdb_id": "nm0000093", // optional, "" if unknown
|
||||
// "tmdb_id": "287", // optional, "" if unknown
|
||||
// "jellyfin_id": "abc123-guid", // optional, "" unless from make_jellyfin_gallery.py
|
||||
// "name": "Brad Pitt",
|
||||
// "source_images": ["img1.jpg", "img2.jpg"],
|
||||
// "embeddings": [[0.012, -0.034, ...], ...] // one 512-float array per image
|
||||
|
||||
@@ -37,6 +37,7 @@ struct IdentityMatcherFunc {
|
||||
, ratio_(cfg.match_ratio)
|
||||
, ratio_ceil_(cfg.match_ratio_ceil)
|
||||
{
|
||||
std::cerr << "[identity_matcher] flattening gallery embeddings...\n";
|
||||
for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
|
||||
for (const auto& emb : gallery_.actors[ai].embeddings) {
|
||||
flat_emb_.push_back(emb);
|
||||
@@ -44,7 +45,10 @@ struct IdentityMatcherFunc {
|
||||
}
|
||||
}
|
||||
|
||||
cal_ = calibrate_gallery(flat_emb_, flat_actor_);
|
||||
std::cerr << "[identity_matcher] starting calibration ("
|
||||
<< flat_emb_.size() << " embeddings)...\n";
|
||||
cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_,
|
||||
cfg.gallery_path + ".calib_cache.json");
|
||||
|
||||
if (cal_.valid) {
|
||||
std::cerr << "[identity_matcher] calibrated Bayesian matching"
|
||||
@@ -126,9 +130,11 @@ struct IdentityMatcherFunc {
|
||||
ia.track_id = tf.track_ids[fi];
|
||||
|
||||
if (accept) {
|
||||
ia.actor_idx = best_actor;
|
||||
ia.name = gallery_.actors[best_actor].name;
|
||||
ia.imdb_id = gallery_.actors[best_actor].imdb_id;
|
||||
ia.actor_idx = best_actor;
|
||||
ia.name = gallery_.actors[best_actor].name;
|
||||
ia.imdb_id = gallery_.actors[best_actor].imdb_id;
|
||||
ia.tmdb_id = gallery_.actors[best_actor].tmdb_id;
|
||||
ia.jellyfin_id = gallery_.actors[best_actor].jellyfin_id;
|
||||
ia.similarity = cal_.valid
|
||||
? cal_.probability(best_s, log_prior_odds_)
|
||||
: best_s;
|
||||
|
||||
@@ -18,7 +18,16 @@ using json = nlohmann::json;
|
||||
// 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], ...] }] }
|
||||
// Output: {
|
||||
// "schema_version": 1, "movie": "...", "sample_fps": ..., "anneal_sec": ...,
|
||||
// "actors": [{ "name", "imdb_id", "tmdb_id", "jellyfin_id", "scenes": [[t0,t1], ...] }]
|
||||
// }
|
||||
// This is the spec consumed by the Jellyfin plugin: each actor carries every
|
||||
// identity key the gallery knows (empty string if not resolved). The plugin
|
||||
// should prefer "jellyfin_id" (direct Person item GUID) when non-empty, and
|
||||
// otherwise resolve "imdb_id"/"tmdb_id" against the item's People ProviderIds.
|
||||
// To find who's on screen at timestamp t, scan each actor's "scenes" for a
|
||||
// window where start <= t <= end.
|
||||
//
|
||||
// Verbosity::standard — per-frame detail including bboxes, similarity, unknowns.
|
||||
// Output: { "frames": [{ "t", "identified": [...], "unknowns": [...] }] }
|
||||
@@ -55,6 +64,10 @@ struct ResultSinkFunc {
|
||||
}
|
||||
|
||||
private:
|
||||
// Bump when the minimal/standard output JSON structure changes in a way
|
||||
// the Jellyfin plugin needs to detect.
|
||||
static constexpr int kSchemaVersion = 1;
|
||||
|
||||
static int count_known(const std::vector<IdentifiedActor>& v) {
|
||||
int n = 0;
|
||||
for (const auto& a : v) if (a.actor_idx >= 0) ++n;
|
||||
@@ -73,6 +86,7 @@ private:
|
||||
if (cfg_.verbosity == Verbosity::xray) {
|
||||
root = build_xray();
|
||||
} else {
|
||||
root["schema_version"] = kSchemaVersion;
|
||||
root["movie"] = cfg_.movie_path;
|
||||
root["sample_fps"] = cfg_.sample_fps;
|
||||
root["anneal_sec"] = cfg_.anneal_sec;
|
||||
@@ -91,20 +105,20 @@ private:
|
||||
}
|
||||
|
||||
struct ActorWindow {
|
||||
std::string name, imdb_id;
|
||||
std::string name, imdb_id, tmdb_id, jellyfin_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; };
|
||||
struct Info { std::string name, imdb_id, tmdb_id, jellyfin_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};
|
||||
actor_info[ia.actor_idx] = {ia.name, ia.imdb_id, ia.tmdb_id, ia.jellyfin_id};
|
||||
timestamps[ia.actor_idx].push_back(frame.timestamp_sec);
|
||||
}
|
||||
}
|
||||
@@ -112,8 +126,10 @@ private:
|
||||
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;
|
||||
aw.name = actor_info[idx].name;
|
||||
aw.imdb_id = actor_info[idx].imdb_id;
|
||||
aw.tmdb_id = actor_info[idx].tmdb_id;
|
||||
aw.jellyfin_id = actor_info[idx].jellyfin_id;
|
||||
|
||||
double win_start = ts_vec[0], win_end = ts_vec[0];
|
||||
for (size_t i = 1; i < ts_vec.size(); ++i) {
|
||||
@@ -136,9 +152,11 @@ private:
|
||||
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);
|
||||
ja["name"] = aw.name;
|
||||
ja["imdb_id"] = aw.imdb_id;
|
||||
ja["tmdb_id"] = aw.tmdb_id;
|
||||
ja["jellyfin_id"] = aw.jellyfin_id;
|
||||
ja["scenes"] = std::move(windows);
|
||||
actors.push_back(std::move(ja));
|
||||
}
|
||||
return actors;
|
||||
@@ -178,8 +196,10 @@ private:
|
||||
|
||||
if (ia.actor_idx >= 0) {
|
||||
json ja;
|
||||
ja["name"] = ia.name;
|
||||
ja["imdb_id"] = ia.imdb_id;
|
||||
ja["name"] = ia.name;
|
||||
ja["imdb_id"] = ia.imdb_id;
|
||||
ja["tmdb_id"] = ia.tmdb_id;
|
||||
ja["jellyfin_id"] = ia.jellyfin_id;
|
||||
ja["similarity"] = ia.similarity;
|
||||
ja["track_id"] = ia.track_id;
|
||||
ja["bbox"] = jbox;
|
||||
|
||||
@@ -41,6 +41,8 @@ struct SceneTrackerFunc {
|
||||
slot.last_crop = ia.crop;
|
||||
slot.name = ia.name;
|
||||
slot.imdb_id = ia.imdb_id;
|
||||
slot.tmdb_id = ia.tmdb_id;
|
||||
slot.jellyfin_id = ia.jellyfin_id;
|
||||
// Keep the best (highest) similarity seen in this window
|
||||
if (ia.similarity > slot.best_similarity)
|
||||
slot.best_similarity = ia.similarity;
|
||||
@@ -60,9 +62,11 @@ struct SceneTrackerFunc {
|
||||
|
||||
for (const auto& [actor_idx, slot] : active_) {
|
||||
IdentifiedActor ia;
|
||||
ia.actor_idx = actor_idx;
|
||||
ia.name = slot.name;
|
||||
ia.imdb_id = slot.imdb_id;
|
||||
ia.actor_idx = actor_idx;
|
||||
ia.name = slot.name;
|
||||
ia.imdb_id = slot.imdb_id;
|
||||
ia.tmdb_id = slot.tmdb_id;
|
||||
ia.jellyfin_id = slot.jellyfin_id;
|
||||
ia.similarity = slot.best_similarity;
|
||||
ia.bbox = slot.last_bbox;
|
||||
ia.crop = slot.last_crop;
|
||||
@@ -85,6 +89,8 @@ private:
|
||||
cv::Mat last_crop;
|
||||
std::string name;
|
||||
std::string imdb_id;
|
||||
std::string tmdb_id;
|
||||
std::string jellyfin_id;
|
||||
};
|
||||
|
||||
double extinction_sec_;
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
// sae_embed — Python module wrapping FaceEmbedderEngine (SCRFD + ArcFace).
|
||||
//
|
||||
// Loads both ONNX sessions once per FaceEmbedder instance, then embeds many
|
||||
// images via repeated embed() calls — avoiding the per-process model-load
|
||||
// cost of the embed_faces CLI when embedding a large gallery.
|
||||
|
||||
#include "face_embedder_engine.hpp"
|
||||
|
||||
#include <nanobind/nanobind.h>
|
||||
#include <nanobind/stl/optional.h>
|
||||
#include <nanobind/stl/string.h>
|
||||
#include <nanobind/stl/vector.h>
|
||||
|
||||
namespace nb = nanobind;
|
||||
using namespace nb::literals;
|
||||
|
||||
NB_MODULE(sae_embed, m) {
|
||||
m.doc() = "SCRFD + ArcFace face embedding, models loaded once per FaceEmbedder";
|
||||
|
||||
nb::class_<FaceEmbedResult>(m, "FaceResult")
|
||||
.def_ro("ok", &FaceEmbedResult::ok)
|
||||
.def_ro("error", &FaceEmbedResult::error)
|
||||
.def_ro("confidence", &FaceEmbedResult::confidence)
|
||||
.def_prop_ro("embedding", [](const FaceEmbedResult& r) -> std::optional<std::vector<float>> {
|
||||
if (!r.ok) return std::nullopt;
|
||||
return std::vector<float>(r.embedding.begin(), r.embedding.end());
|
||||
})
|
||||
.def_prop_ro("bbox", [](const FaceEmbedResult& r) {
|
||||
return std::vector<float>{r.bbox[0], r.bbox[1], r.bbox[2], r.bbox[3]};
|
||||
});
|
||||
|
||||
nb::class_<FaceEmbedderEngine>(m, "FaceEmbedder")
|
||||
.def(nb::init<std::string, std::string, float, float, int>(),
|
||||
"detector_model"_a, "arcface_model"_a,
|
||||
"conf"_a = 0.5f, "nms"_a = 0.4f, "max_side"_a = 500)
|
||||
.def("embed", &FaceEmbedderEngine::embed_path, "path"_a,
|
||||
nb::call_guard<nb::gil_scoped_release>(),
|
||||
"Detect the highest-confidence face in the image, align it, and "
|
||||
"return a FaceResult with its 512-d ArcFace embedding.");
|
||||
}
|
||||
@@ -92,6 +92,8 @@ struct IdentifiedActor {
|
||||
int track_id{-1}; // face track ID from FaceTrackerFunc
|
||||
std::string name;
|
||||
std::string imdb_id;
|
||||
std::string tmdb_id;
|
||||
std::string jellyfin_id; // Jellyfin Person item GUID, if gallery was built from Jellyfin
|
||||
float similarity{0.f}; // calibrated P(match) or cosine similarity; 0 for unknowns
|
||||
cv::Rect2f bbox;
|
||||
cv::Mat crop; // 112×112 aligned crop (stored as shared_ptr by KPN)
|
||||
@@ -116,6 +118,8 @@ struct SceneAnnotation {
|
||||
struct ActorGallery {
|
||||
struct Actor {
|
||||
std::string imdb_id;
|
||||
std::string tmdb_id;
|
||||
std::string jellyfin_id; // Jellyfin Person item GUID, if known
|
||||
std::string name;
|
||||
std::vector<Embedding> embeddings; // one per reference image
|
||||
std::vector<std::string> source_images;
|
||||
|
||||
Reference in New Issue
Block a user