src/embed_face.py

"""
Face embedding service — extracted from ai_manager.py.
Handles: CLAHE enhancement, multi-scale SCRFD detection, ArcFace + AdaFace embedding,
IOU deduplication, face crop thumbnails.
"""
import base64
import io
import os
import sys
import threading

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

# ── Config (env-tunable) ─────────────────────────────────────────────────────
DET_SIZE_PRIMARY   = (640, 640)
DET_SCALES         = [(1280, 1280), (960, 960), (640, 640)]
IOU_DEDUP_THRESHOLD = float(os.getenv("IOU_DEDUP_THRESHOLD", "0.4"))
MIN_FACE_SIZE      = int(os.getenv("MIN_FACE_SIZE", "30"))
MAX_FACES_PER_IMAGE = int(os.getenv("MAX_FACES_PER_IMAGE", "20"))
FACE_CROP_THUMB_SIZE = int(os.getenv("FACE_CROP_THUMB_SIZE", "112"))
FACE_CROP_QUALITY  = int(os.getenv("FACE_CROP_QUALITY", "85"))
FACE_CROP_PADDING  = float(os.getenv("FACE_CROP_PADDING", "0.2"))
ADAFACE_CROP_PADDING = float(os.getenv("ADAFACE_CROP_PADDING", "0.1"))
ADAFACE_DIM        = 512
ENABLE_ADAFACE     = int(os.getenv("ENABLE_ADAFACE", "1"))
HF_TOKEN           = os.getenv("HF_TOKEN", "")
FAST_DETECT        = bool(int(os.getenv("FAST_DETECT", "1")))


# ── Utility functions ─────────────────────────────────────────────────────────

def _clahe_enhance(bgr: np.ndarray) -> np.ndarray:
    lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB)
    l_ch, a_ch, b_ch = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l_eq = clahe.apply(l_ch)
    return cv2.cvtColor(cv2.merge([l_eq, a_ch, b_ch]), cv2.COLOR_LAB2BGR)


def _iou(box_a: list, box_b: list) -> float:
    xa, ya = max(box_a[0], box_b[0]), max(box_a[1], box_b[1])
    xb, yb = min(box_a[2], box_b[2]), min(box_a[3], box_b[3])
    inter = max(0, xb - xa) * max(0, yb - ya)
    if inter == 0:
        return 0.0
    area_a = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
    area_b = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
    return inter / (area_a + area_b - inter)


def _dedup_faces(faces_list: list, iou_thresh: float = IOU_DEDUP_THRESHOLD) -> list:
    if not faces_list:
        return []
    faces_list = sorted(faces_list, key=lambda f: float(f.det_score), reverse=True)
    kept = []
    for face in faces_list:
        b = face.bbox.astype(int)
        box = [b[0], b[1], b[2], b[3]]
        if not any(_iou(box, [k.bbox.astype(int)[i] for i in range(4)]) > iou_thresh for k in kept):
            kept.append(face)
    return kept


def _crop_to_b64(img_bgr: np.ndarray, x1: int, y1: int, x2: int, y2: int) -> str:
    H, W = img_bgr.shape[:2]
    w, h = x2 - x1, y2 - y1
    pad_x, pad_y = int(w * FACE_CROP_PADDING), int(h * FACE_CROP_PADDING)
    cx1, cy1 = max(0, x1 - pad_x), max(0, y1 - pad_y)
    cx2, cy2 = min(W, x2 + pad_x), min(H, y2 + pad_y)
    crop = img_bgr[cy1:cy2, cx1:cx2]
    if crop.size == 0:
        return ""
    pil = Image.fromarray(crop[:, :, ::-1]).resize((FACE_CROP_THUMB_SIZE, FACE_CROP_THUMB_SIZE), Image.LANCZOS)
    buf = io.BytesIO()
    pil.save(buf, format="JPEG", quality=FACE_CROP_QUALITY)
    return base64.b64encode(buf.getvalue()).decode()


def _face_crop_for_adaface(img_bgr: np.ndarray, x1: int, y1: int, x2: int, y2: int):
    H, W = img_bgr.shape[:2]
    w, h = x2 - x1, y2 - y1
    pad_x, pad_y = int(w * ADAFACE_CROP_PADDING), int(h * ADAFACE_CROP_PADDING)
    cx1, cy1 = max(0, x1 - pad_x), max(0, y1 - pad_y)
    cx2, cy2 = min(W, x2 + pad_x), min(H, y2 + pad_y)
    crop = img_bgr[cy1:cy2, cx1:cx2]
    if crop.size == 0:
        return None
    rgb = crop[:, :, ::-1].copy()
    pil = Image.fromarray(rgb).resize((112, 112), Image.LANCZOS)
    arr = np.array(pil, dtype=np.float32) / 255.0
    arr = (arr - 0.5) / 0.5
    return arr.transpose(2, 0, 1)


# ── FaceEmbedder class ────────────────────────────────────────────────────────

class FaceEmbedder:
    def __init__(self):
        from insightface.app import FaceAnalysis
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if self.device == "cuda" else ["CPUExecutionProvider"]
        self.face_app = FaceAnalysis(name="buffalo_l", providers=providers)
        self.face_app.prepare(ctx_id=0 if self.device == "cuda" else -1, det_size=DET_SIZE_PRIMARY)
        self.face_app.get(np.zeros((112, 112, 3), dtype=np.uint8))  # warm-up
        self._face_lock = threading.Lock()
        self.adaface_model = None
        self._load_adaface()

    def _load_adaface(self) -> None:
        if not ENABLE_ADAFACE:
            return
        REPO_ID = "minchul/cvlface_adaface_ir50_ms1mv2"
        CACHE_PATH = os.path.expanduser("~/.cvlface_cache/minchul/cvlface_adaface_ir50_ms1mv2")
        try:
            from huggingface_hub import hf_hub_download
            from transformers import AutoModel as _HFAutoModel
            os.makedirs(CACHE_PATH, exist_ok=True)
            hf_hub_download(repo_id=REPO_ID, filename="files.txt", token=HF_TOKEN, local_dir=CACHE_PATH, local_dir_use_symlinks=False)
            with open(os.path.join(CACHE_PATH, "files.txt")) as f:
                extra = [x.strip() for x in f.read().split("\n") if x.strip()]
            for fname in extra + ["config.json", "wrapper.py", "model.safetensors"]:
                if not os.path.exists(os.path.join(CACHE_PATH, fname)):
                    hf_hub_download(repo_id=REPO_ID, filename=fname, token=HF_TOKEN, local_dir=CACHE_PATH, local_dir_use_symlinks=False)
            cwd = os.getcwd()
            os.chdir(CACHE_PATH)
            sys.path.insert(0, CACHE_PATH)
            try:
                model = _HFAutoModel.from_pretrained(CACHE_PATH, trust_remote_code=True, token=HF_TOKEN)
            finally:
                os.chdir(cwd)
                if CACHE_PATH in sys.path:
                    sys.path.remove(CACHE_PATH)
            self.adaface_model = model.to(self.device).eval()
        except Exception:
            self.adaface_model = None

    def _adaface_embed_single(self, face_arr_chw) -> np.ndarray | None:
        if self.adaface_model is None or face_arr_chw is None:
            return None
        try:
            t = torch.from_numpy(face_arr_chw).unsqueeze(0).to(self.device)
            if self.device == "cuda":
                t = t.half()
            with torch.no_grad():
                out = self.adaface_model(t)
            emb = out if isinstance(out, torch.Tensor) else out.embedding
            return F.normalize(emb.float(), p=2, dim=1)[0].cpu().numpy()
        except Exception:
            return None

    def _adaface_embed_batch(self, face_arr_list: list) -> list:
        if self.adaface_model is None or not face_arr_list:
            return [None] * len(face_arr_list)
        valid_indices = [i for i, arr in enumerate(face_arr_list) if arr is not None]
        if not valid_indices:
            return [None] * len(face_arr_list)
        try:
            batch = np.stack([face_arr_list[i] for i in valid_indices])
            t = torch.from_numpy(batch).to(self.device)
            if self.device == "cuda":
                t = t.half()
            with torch.no_grad():
                out = self.adaface_model(t)
            emb = out if isinstance(out, torch.Tensor) else out.embedding
            normed = F.normalize(emb.float(), p=2, dim=1).cpu().numpy()
            results = [None] * len(face_arr_list)
            for batch_i, orig_i in enumerate(valid_indices):
                results[orig_i] = normed[batch_i]
            return results
        except Exception:
            return [self._adaface_embed_single(arr) for arr in face_arr_list]

    def embed(self, image_bytes: bytes, quality_gate: float = 0.35) -> list[dict]:
        try:
            pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            img_np = np.array(pil)
            if img_np.dtype != np.uint8:
                img_np = (img_np * 255).astype(np.uint8)
            bgr = img_np[:, :, ::-1].copy()
            bgr_enhanced = _clahe_enhance(bgr)
            H, W = bgr.shape[:2]

            all_raw_faces = []
            for scale in DET_SCALES:
                scale_w, scale_h = min(W, scale[0]), min(H, scale[1])
                bgr_scaled = bgr_enhanced if scale_w == W and scale_h == H else cv2.resize(bgr_enhanced, (scale_w, scale_h))
                try:
                    self.face_app.det_model.input_size = scale
                    with self._face_lock:
                        faces_at_scale = self.face_app.get(bgr_scaled)
                    sx, sy = W / scale_w, H / scale_h
                    for f in faces_at_scale:
                        if sx != 1.0 or sy != 1.0:
                            f.bbox[0] *= sx; f.bbox[1] *= sy; f.bbox[2] *= sx; f.bbox[3] *= sy
                    all_raw_faces.extend(faces_at_scale)
                except Exception:
                    pass
                if FAST_DETECT and len(all_raw_faces) >= MAX_FACES_PER_IMAGE:
                    break

            bgr_flip = cv2.flip(bgr_enhanced, 1)
            try:
                self.face_app.det_model.input_size = DET_SIZE_PRIMARY
                with self._face_lock:
                    faces_flip = self.face_app.get(bgr_flip)
                for f in faces_flip:
                    x1, y1, x2, y2 = f.bbox
                    f.bbox[0], f.bbox[2] = W - x2, W - x1
                all_raw_faces.extend(faces_flip)
            except Exception:
                pass

            self.face_app.det_model.input_size = DET_SIZE_PRIMARY
            faces = _dedup_faces(all_raw_faces)

            # Pass 1: validate and collect crops
            valid_faces = []
            for face in faces:
                if len(valid_faces) >= MAX_FACES_PER_IMAGE:
                    break
                bbox_raw = face.bbox.astype(int)
                x1, y1, x2, y2 = bbox_raw
                x1, y1 = max(0, x1), max(0, y1)
                x2, y2 = min(bgr.shape[1], x2), min(bgr.shape[0], y2)
                w, h = x2 - x1, y2 - y1
                if w < MIN_FACE_SIZE or h < MIN_FACE_SIZE:
                    continue
                det_score = float(face.det_score) if hasattr(face, "det_score") else 1.0
                if det_score < quality_gate or face.embedding is None:
                    continue
                arcface_vec = face.embedding.astype(np.float32)
                n = np.linalg.norm(arcface_vec)
                arcface_vec = arcface_vec / n if n > 0 else arcface_vec
                face_chw = _face_crop_for_adaface(bgr, x1, y1, x2, y2)
                valid_faces.append({
                    "x1": x1, "y1": y1, "x2": x2, "y2": y2,
                    "w": w, "h": h, "det_score": det_score,
                    "arcface_vec": arcface_vec, "face_chw": face_chw,
                })

            # Pass 2: batch AdaFace
            adaface_vecs = self._adaface_embed_batch([f["face_chw"] for f in valid_faces])

            # Pass 3: assemble fused vectors
            results = []
            for idx, fd in enumerate(valid_faces):
                adaface_vec = adaface_vecs[idx]
                arcface_vec = fd["arcface_vec"]
                if adaface_vec is not None:
                    fused_raw = np.concatenate([arcface_vec, adaface_vec])
                    vec_mode = "fused"
                else:
                    fused_raw = np.concatenate([arcface_vec, arcface_vec.copy()])
                    vec_mode = "arcface_mirror"
                n2 = np.linalg.norm(fused_raw)
                final_vec = (fused_raw / n2) if n2 > 0 else fused_raw
                results.append({
                    "type": "face",
                    "vector": final_vec.tolist(),
                    "face_idx": idx,
                    "bbox": [int(fd["x1"]), int(fd["y1"]), int(fd["w"]), int(fd["h"])],
                    "face_crop": _crop_to_b64(bgr, fd["x1"], fd["y1"], fd["x2"], fd["y2"]),
                    "det_score": fd["det_score"],
                    "face_width_px": int(fd["w"]),
                    "vec_mode": vec_mode,
                })
            return results
        except Exception:
            return []