models/sprite_processor.py

import os
from concurrent.futures import ThreadPoolExecutor, as_completed

import numpy as np

from models.face_recognition import (
    AdaFaceRecognition,
    extract_faces,
    _get_face_analyzer,
    _dist_to_confidence,
    _face_to_base64,
)
from utils.vtt_parser import parse_sprite_vtt

HIT_CONF_THRESH = 0.15
CAST_SCORE_THRESH = 0.9
CAST_HIT_THRESH = 2
RECLAIM_SIM_THRESH = 0.35
GAP_TOLERANCE = 2
TOP_K_CANDIDATES = 10
CONFIDENCE_FLOOR = 0.05

# Core-aware parallelism. On Hugging Face Spaces (low vCPU), this auto-scales down.
MAX_WORKERS = max(1, os.cpu_count() or 1)

# Shared buffalo_sc 640x640 analyzer — ONNX sessions are thread-safe.
_sprite_ana = None


def _get_sprite_analyzer():
    global _sprite_ana
    if _sprite_ana is None:
        _sprite_ana = _get_face_analyzer(det_size=(640, 640))
    return _sprite_ana


def _detect_one(frame_args):
    """Detect faces in a single frame. Called by ThreadPoolExecutor."""
    idx, frame, mc, ana = frame_args
    results = []
    for d in extract_faces(frame, min_confidence=mc, analyzer=ana):
        results.append((d['face'], idx))
    return results


def _detect_faces(frames, min_confidence):
    """Per-frame face detection with buffalo_sc at 640x640.

    Single-core: sequential per-frame.
    Multi-core: parallel per-frame via ThreadPoolExecutor.
    """
    n_frames = len(frames)
    if n_frames == 0:
        return []

    ana = _get_sprite_analyzer()

    if MAX_WORKERS <= 1:
        results = []
        for idx, frame in enumerate(frames):
            for d in extract_faces(frame, min_confidence=min_confidence, analyzer=ana):
                results.append((d['face'], idx))
        return results

    face_crops = []
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
        args = [(idx, f, min_confidence, ana) for idx, f in enumerate(frames)]
        for fut in as_completed(pool.submit(_detect_one, a) for a in args):
            face_crops.extend(fut.result())

    face_crops.sort(key=lambda x: x[1])
    return face_crops


def identify_sprite_performers(sprite_image, vtt_text, data_manager,
                                detection_threshold=0.5):
    yield (0.0, "Parsing VTT cue sheet")
    cues = parse_sprite_vtt(vtt_text)
    if not cues:
        raise ValueError("No cues found in VTT file")

    yield (0.05, "Extracting sprite frames")
    sprite = np.array(sprite_image)
    frames = [sprite[c['y']:c['y'] + c['h'], c['x']:c['x'] + c['w']] for c in cues]

    yield (0.1, f"Detecting faces across {len(frames)} frames\u2026")
    face_crops = _detect_faces(frames, min_confidence=detection_threshold)
    yield (0.6, f"Found faces in {len(set(fc[1] for fc in face_crops))} frames")

    if not face_crops:
        raise ValueError("Face detection found no faces in any sprite frame.")

    yield (0.6, f"Generating embeddings for {len(face_crops)} faces\u2026")
    face_embs = AdaFaceRecognition().get_adaface_embeddings_batch(
        np.stack([fc[0] for fc in face_crops], axis=0)
    )

    yield (0.7, f"Matching {len(face_embs)} faces against performer index\u2026")
    performer_data = {}
    face_results = []
    for emb in face_embs:
        ids, dists = data_manager.query_adaface_index(emb, TOP_K_CANDIDATES)
        confident = []
        all_hits = []
        for pid, dist in zip(ids, dists):
            conf = _dist_to_confidence(dist)
            if conf < CONFIDENCE_FLOOR:
                break
            info = data_manager.get_performer_info(pid, conf, dist)
            if info and info.get('name'):
                p = performer_data.setdefault(pid, {"score": 0.0, "hits": 0, "conf_hits": 0})
                p["score"] += conf
                p["hits"] += 1
                all_hits.append((pid, conf))
                if conf >= HIT_CONF_THRESH:
                    p["conf_hits"] += 1
                    confident.append(pid)
        face_results.append((confident, emb, all_hits))

    yield (0.85, "Clustering results by performer")
    cast = [pid for pid, p in performer_data.items()
            if p["score"] >= CAST_SCORE_THRESH and p["conf_hits"] >= CAST_HIT_THRESH]

    if not cast and face_crops:
        raise ValueError(
            f"Faces detected in {len(set(fc[1] for fc in face_crops))} frames, but "
            "no performer could be identified with confidence. "
            "Individual frame matches were too weak to cluster into a reliable result."
        )

    centroids = {}
    for pid in cast:
        embs = [face_results[i][1] for i in range(len(face_results)) if pid in face_results[i][0]]
        if embs:
            centroids[pid] = np.mean(embs, axis=0)

    cue_sets = {pid: set() for pid in cast}
    face_indices = {pid: [] for pid in cast}
    for i, (confident, _, _) in enumerate(face_results):
        for pid in confident:
            if pid not in cue_sets:
                continue
            cue_sets[pid].add(face_crops[i][1])
            face_indices[pid].append(i)

    for i, (confident, emb, _) in enumerate(face_results):
        if not confident and centroids:
            best = max(((pid, float(np.dot(emb, c))) for pid, c in centroids.items()), key=lambda x: x[1], default=(None, 0))
            if best[0] is not None and best[1] >= RECLAIM_SIM_THRESH:
                cue_sets[best[0]].add(face_crops[i][1])
                face_indices[best[0]].append(i)

    yield (0.95, "Building response")
    response = []
    for pid in cast:
        info = data_manager.get_performer_info(pid, 0, 0)
        if not info:
            continue

        centroid = centroids[pid]
        best_idx = max(face_indices[pid], key=lambda i: float(np.dot(face_results[i][1], centroid))) if face_indices[pid] else None
        thumb = _face_to_base64(face_crops[best_idx][0]) if best_idx is not None else info.get('image', '')

        shot_groups = _group_adjacent_indices(cue_sets[pid])
        shot_confs = []
        for group in shot_groups:
            max_in_group = 0.0
            for idx in face_indices[pid]:
                if face_crops[idx][1] not in group:
                    continue
                for hit_pid, conf in face_results[idx][2]:
                    if hit_pid == pid and conf > max_in_group:
                        max_in_group = conf
            shot_confs.append(max_in_group)

        response.append({
            "id": pid,
            "name": info['name'],
            "confidence": int((1.0 - np.prod([1.0 - c for c in shot_confs])) * 100),
            "thumbnail": thumb,
            "time_ranges": _merge_time_ranges(cue_sets[pid], cues),
            "hit_count": performer_data[pid]["hits"],
            "shot_count": len(shot_confs),
            "performer_url": info.get('performer_url', ''),
            "source": info.get('source', ''),
            "gender": info.get('gender'),
        })

    response.sort(key=lambda x: x['hit_count'], reverse=True)

    yield {"performers": response}


def _group_adjacent_indices(indices, gap_tolerance=GAP_TOLERANCE):
    if not indices:
        return []
    sorted_idx = sorted(set(indices))
    groups = [[sorted_idx[0]]]
    for i in range(1, len(sorted_idx)):
        if sorted_idx[i] - sorted_idx[i - 1] <= gap_tolerance + 1:
            groups[-1].append(sorted_idx[i])
        else:
            groups.append([sorted_idx[i]])
    return groups


def _merge_time_ranges(cue_indices, cues, gap_tolerance=GAP_TOLERANCE):
    groups = _group_adjacent_indices(cue_indices, gap_tolerance)
    return [{"start": cues[g[0]]["start"], "end": cues[g[-1]]["end"]} for g in groups]