cores/face/helpers.py

"""Face helpers — box conversions, embedding distance, gallery matching."""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np

from cores.vision.geometry import BBox, crop_region


# --------------------------------------------------------------------------- #
# Box format conversions
# --------------------------------------------------------------------------- #
def xywh_to_xyxy(x: int, y: int, w: int, h: int) -> Tuple[int, int, int, int]:
    """(x, y, w, h) -> (x1, y1, x2, y2)."""
    return (x, y, x + w, y + h)


def xyxy_to_xywh(x1: int, y1: int, x2: int, y2: int) -> Tuple[int, int, int, int]:
    """(x1, y1, x2, y2) -> (x, y, w, h)."""
    return (x1, y1, x2 - x1, y2 - y1)


def xywh_to_face_recognition_tuple(x: int, y: int, w: int, h: int) -> Tuple[int, int, int, int]:
    """Convert (x, y, w, h) to (top, right, bottom, left) used by face_recognition."""
    return (y, x + w, y + h, x)


# --------------------------------------------------------------------------- #
# Embedding distance / similarity
# --------------------------------------------------------------------------- #
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity between two 1-D vectors.  Returns float in [-1, 1]."""
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na == 0 or nb == 0:
        return 0.0
    return float(np.dot(a, b) / (na * nb))


def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    """Euclidean distance between two 1-D vectors."""
    return float(np.linalg.norm(a - b))


# --------------------------------------------------------------------------- #
# Gallery matching
# --------------------------------------------------------------------------- #
def best_match(
    query: np.ndarray,
    gallery: Dict[str, List[np.ndarray]],
    metric: str = "cosine",
) -> Tuple[Optional[str], float, Dict[str, float]]:
    """Find the best matching person in the gallery for a query embedding.

    Args:
        query: 1-D embedding vector.
        gallery: dict mapping person_name -> list of reference embeddings.
        metric: "cosine" (higher = better) or "euclidean" (lower = better).

    Returns:
        (best_name, best_score, all_scores)
        - For cosine: best_score is the highest similarity.
        - For euclidean: best_score is the smallest distance.
        - best_name is None if the gallery is empty.
    """
    if not gallery:
        return None, 0.0, {}

    all_scores: Dict[str, float] = {}
    best_name: Optional[str] = None
    best_score: float = -1.0 if metric == "cosine" else float("inf")

    for name, embeddings in gallery.items():
        if not embeddings:
            continue
        if metric == "cosine":
            scores = [cosine_similarity(query, ref) for ref in embeddings]
            score = max(scores)  # higher = better
        else:
            scores = [euclidean_distance(query, ref) for ref in embeddings]
            score = min(scores)  # lower = better
        all_scores[name] = round(score, 4)
        if (metric == "cosine" and score > best_score) or \
           (metric == "euclidean" and score < best_score):
            best_score = score
            best_name = name

    return best_name, round(best_score, 4), all_scores


# --------------------------------------------------------------------------- #
# Face-crop extraction
# --------------------------------------------------------------------------- #
def extract_face_crops(
    img: np.ndarray,
    boxes: List[dict],
    margin: float = 0.2,
) -> List[np.ndarray]:
    """Extract face crops from an image given a list of box dicts.

    Each box dict must have keys: x, y, w, h.
    """
    crops: List[np.ndarray] = []
    for b in boxes:
        bbox = BBox(b["x"], b["y"], b["w"], b["h"])
        crops.append(crop_region(img, bbox, margin=margin))
    return crops