scene-actor-extraction/src/face_embedder_engine.hpp

#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 "config.hpp"
#include "face_utils.hpp"
#include "inference/face_detector.hpp"
#include "inference/face_embedder.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:
    // detector_engine/arcface_engine are optional paths to pre-built TensorRT
    // engines. They are required when built with SAE_INFERENCE_BACKEND=TRT
    // (which cannot load .onnx directly) and ignored by the ORT backend.
    FaceEmbedderEngine(const std::string& detector_model,
                        const std::string& arcface_model,
                        float conf = 0.5f, float nms = 0.4f, int max_side = 500,
                        const std::string& detector_engine = "",
                        const std::string& arcface_engine = "")
        : max_side_(max_side)
    {
        Config cfg;
        cfg.detector_model  = detector_model;
        cfg.arcface_model   = arcface_model;
        cfg.detector_engine = detector_engine;
        cfg.arcface_engine  = arcface_engine;
        cfg.detector_conf   = conf;
        cfg.detector_nms    = nms;
        detector_ = make_face_detector(cfg);
        embedder_ = make_face_embedder(cfg);
    }

    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()) {
            cv::Mat enhanced = enhance_for_retry(img);
            faces = detector_->detect(enhanced);
            if (!faces.empty())
                img = enhanced;
        }
        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;
    }

    // ── Stage accessors ──────────────────────────────────────────────────────
    // embed_mat() above is the whole detect→align→embed chain, which is the
    // right entry point for embedding a gallery image. Studies that need to
    // intervene between the stages — swapping the landmark source, degrading a
    // crop before it reaches the embedder — drive these instead, so they still
    // exercise the shipped detector, alignment and embedder rather than a
    // re-implementation of them.
    std::vector<DetectedFace> detect(const cv::Mat& img) { return detector_->detect(img); }

    Embedding embed_crop(const cv::Mat& crop) { return embedder_->embed_one(crop); }

    // Batched form. A study embedding thousands of crops one at a time pays the
    // per-call overhead thousands of times over; the backend already batches.
    std::vector<Embedding> embed_crops(const std::vector<cv::Mat>& crops) {
        return embedder_->embed(crops);
    }

    int max_batch() const { return embedder_->max_batch(); }

private:
    std::unique_ptr<IFaceDetector>   detector_;
    std::unique_ptr<IFaceEmbedder>   embedder_;
    int                               max_side_;
};