scene-actor-extraction/src/gallery/gallery_builder.cpp

#include "gallery_builder.hpp"
#include "config.hpp"
#include "embedder_stamp.hpp"
#include "face_utils.hpp"
#include "inference/face_detector.hpp"
#include "inference/face_embedder.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) {
    Config icfg;
    icfg.detector_model = cfg.detector_model;
    icfg.arcface_model  = cfg.arcface_model;
    icfg.detector_conf  = cfg.detector_conf;
    icfg.detector_nms   = cfg.detector_nms;
    auto decoder = make_face_detector(icfg);
    auto arcface = make_face_embedder(icfg);

    ActorGallery gallery;

    /// TRACES: GR-004 | SR-001
    // Stamp before the first embedding exists, so there is no window in which a
    // gallery holds vectors without recording what produced them.
    gallery.embedder = make_embedder_stamp(cfg.arcface_model);
    std::cerr << "[build_gallery] embedder: " << gallery.embedder.describe() << "\n";

    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;
}