faster calibration curve generation

jellyfin intergration
2026-06-12 17:54:23 +02:00
parent d753062c6c
commit a1d6759abc
17 changed files with 1379 additions and 166 deletions
@@ -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_;
+};