model.py

"""Face Recognition: who is in this photo, and are two photos the same person?

Single source of truth for the pipeline used by the training code (`training/src/`), the
Hugging Face Space (`space/app.py`), the Inference Endpoint (`handler.py`) and anyone who
downloads this repo from the Hub:

    1. Our FaceDetector (a single-shot CNN trained in this project, `face_detector.safetensors`)
       finds the largest face and its 5 landmarks (eyes, nose tip, mouth corners), and the face is
       rotated and scaled onto a fixed 128x128 template. When no face is found, the center square
       is used instead.
    2. Face embedders trained in this project (torchvision CNNs + a 512-d projection, trained with
       a CosFace margin loss on the LFW identities) turn the face into an L2-normalised embedding;
       the model is an ensemble of them (see `FaceRecognizer`). No pretrained face network is used:
       every network in this file was trained in this project.
    3. A small MLP head (embedding -> 256 -> 42) names the person. Verification compares two embeddings
       with the cosine distance.

    import model
    predictor = model.load("path/to/this/repo", device="cpu")
    predictor.predict("face.jpg")            # {"George W Bush": 0.99, ...} (all 42 people)
    predictor.verify("a.jpg", "b.jpg")       # {"same_person": True, "distance": 0.31, ...}
    predictor.embed("face.jpg")              # 1024 floats (512 per embedder)
"""
from __future__ import annotations

import base64
import binascii
import io
import os
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import torchvision
from huggingface_hub import PyTorchModelHubMixin
from PIL import Image
from safetensors.torch import load_file
from torch import nn

REPO_ID = "shalev396/face-recognition"
FRAMEWORK = "pytorch"
WEIGHTS = "model.safetensors"             # embedders + identity head (PyTorchModelHubMixin)
MAX_SIDE = 1024   # larger photos are downscaled before detection (speed); LFW photos are 250 px
EMBEDDING_DIM = 512
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def cuda_available() -> bool:
    return torch.cuda.is_available()


# --------------------------------------------------------------------------- inputs
def to_pil(image) -> Image.Image:
    """PIL image, file path, raw bytes, base64 / data-URL string or HxWx3 uint8 array -> RGB PIL."""
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    if isinstance(image, np.ndarray):
        return Image.fromarray(np.asarray(image, dtype=np.uint8)).convert("RGB")
    if isinstance(image, (bytes, bytearray)):
        with Image.open(io.BytesIO(image)) as img:
            return img.convert("RGB")
    if isinstance(image, (str, os.PathLike)):
        text = str(image)
        if len(text) < 1024 and Path(text).is_file():
            with Image.open(text) as img:
                return img.convert("RGB")
        payload = text.split(",", 1)[1] if text.startswith("data:") else text
        try:
            return to_pil(base64.b64decode(payload, validate=True))
        except (binascii.Error, ValueError, OSError) as err:
            raise ValueError("expected an image file path or a base64-encoded image") from err
    raise TypeError(f"unsupported image input: {type(image).__name__}")


def _limit_size(img: Image.Image) -> Image.Image:
    if max(img.size) <= MAX_SIDE:
        return img
    img = img.copy()
    img.thumbnail((MAX_SIDE, MAX_SIDE), Image.BILINEAR)
    return img


def center_crop_face(img: Image.Image, crop_size: int) -> Image.Image:
    """Fallback when no face is found: the center square, resized to crop_size."""
    w, h = img.size
    s = min(w, h)
    left, top = (w - s) // 2, (h - s) // 2
    return img.crop((left, top, left + s, top + s)).resize((crop_size, crop_size), Image.BILINEAR)


def face_to_pil(face: torch.Tensor) -> Image.Image:
    """(3, H, W) face tensor with values 0..255 -> viewable PIL crop."""
    arr = face.detach().float().cpu().round().clamp(0, 255).byte().permute(1, 2, 0).numpy()
    return Image.fromarray(arr)


def cosine_distance(a, b) -> float:
    """1 - cosine similarity of two embeddings: 0 = identical direction, 2 = opposite."""
    a, b = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
    a = a / (np.linalg.norm(a) + 1e-12)
    b = b / (np.linalg.norm(b) + 1e-12)
    return float(1.0 - np.dot(a, b))


# --------------------------------------------------------------------------- networks
def build_backbone(arch: str, pretrained: bool = False) -> tuple[nn.Module, int]:
    """torchvision `arch` without its classifier -> (module, feature size). `pretrained` loads the
    ImageNet weights (training only; the deployed weights come from model.safetensors)."""
    net = getattr(torchvision.models, arch)(weights="DEFAULT" if pretrained else None)
    if hasattr(net, "fc"):                       # resnet*
        features = net.fc.in_features
        net.fc = nn.Identity()
    else:                                        # efficientnet* / mobilenet* / convnext*
        features = net.classifier[-1].in_features
        net.classifier[-1] = nn.Identity()
    return net, features


class FaceEmbedNet(nn.Module):
    """One embedder: torchvision backbone -> Dropout -> Linear -> BatchNorm = `embedding_dim` features,
    reading faces resized to `input_size`."""

    def __init__(self, arch: str, input_size: int, embedding_dim: int = EMBEDDING_DIM):
        super().__init__()
        self.arch, self.input_size = arch, int(input_size)
        self.backbone, features = build_backbone(arch)
        self.neck = nn.Sequential(nn.Dropout(0.2), nn.Linear(features, embedding_dim), nn.BatchNorm1d(embedding_dim))
        self.register_buffer("mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1), persistent=False)
        self.register_buffer("std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1), persistent=False)

    def resize(self, faces: torch.Tensor) -> torch.Tensor:
        """(N, 3, crop, crop) faces 0..255 -> float faces at input_size (still 0..255)."""
        return F.interpolate(faces.float(), size=(self.input_size, self.input_size), mode="bilinear",
                             antialias=True, align_corners=False)

    def normalize(self, x: torch.Tensor) -> torch.Tensor:
        """Faces at input_size, 0..255 -> ImageNet-normalised."""
        return (x / 255.0 - self.mean) / self.std

    def features(self, x: torch.Tensor) -> torch.Tensor:
        """Normalised faces -> raw (unnormalised) embeddings. Used by training."""
        return self.neck(self.backbone(x))

    def embed(self, faces: torch.Tensor) -> torch.Tensor:
        """(N, 3, crop, crop) faces 0..255 -> (N, embedding_dim) L2-normalised, flip-averaged embeddings."""
        x = self.normalize(self.resize(faces))
        return F.normalize(self.features(x) + self.features(torch.flip(x, dims=[3])), dim=1)


class FaceRecognizer(
    nn.Module,
    PyTorchModelHubMixin,
    tags=["ml-lab", "face-recognition", "image-classification"],
    repo_url="https://github.com/shalev396/ml-lab",
    pipeline_tag="image-classification",
    license="mit",
):
    """An ensemble of face embedders + the identity MLP head.

    `members` lists the embedders ({"arch", "input_size"}); each gives a 512-d L2-normalised
    embedding and the face embedding is their concatenation divided by sqrt(len(members)), so it
    is L2-normalised too and the cosine distance averages the members' cosine distances.
    `forward()` maps faces to identity logits through the head. The extra args are runtime
    metadata written to config.json by the mixin: crop size and the verification threshold.
    """

    def __init__(self, class_names: list[str], members: list[dict] | None = None,
                 embedding_dim: int = EMBEDDING_DIM, hidden_units: int = 256, dropout: float = 0.3,
                 crop_size: int = 128, verify_threshold: float = 0.5):
        super().__init__()
        members = members or [{"arch": "resnet18", "input_size": 112}]
        self.class_names = list(class_names)
        self.crop_size = int(crop_size)
        self.verify_threshold = float(verify_threshold)
        self.members = nn.ModuleList(FaceEmbedNet(m["arch"], m["input_size"], embedding_dim) for m in members)
        self.embedding_dim = embedding_dim * len(members)
        self.head = nn.Sequential(nn.Linear(self.embedding_dim, hidden_units), nn.ReLU(), nn.Dropout(dropout),
                                  nn.Linear(hidden_units, len(self.class_names)))

    def set_verify_threshold(self, threshold: float) -> None:
        """Change the default verification threshold, also in the config that save_pretrained writes."""
        self.verify_threshold = float(threshold)
        if isinstance(self._hub_mixin_config, dict):
            self._hub_mixin_config["verify_threshold"] = self.verify_threshold

    def embed(self, faces: torch.Tensor) -> torch.Tensor:
        """(N, 3, crop, crop) faces 0..255 -> (N, 512 * members) L2-normalised embeddings."""
        parts = [member.embed(faces) for member in self.members]
        return torch.cat(parts, dim=1) / len(parts) ** 0.5

    def classify(self, embeddings: torch.Tensor) -> torch.Tensor:
        """(N, embedding_dim) normalised embeddings -> (N, n_classes) logits."""
        return self.head(embeddings)

    def forward(self, faces: torch.Tensor) -> torch.Tensor:
        return self.classify(self.embed(faces))


def _require(path: Path) -> Path:
    if not path.is_file():
        raise FileNotFoundError(
            f"{path.name} not found in {path.parent}: export the model first or download it with "
            f"huggingface_hub.snapshot_download('{REPO_ID}')."
        )
    return path


# --------------------------------------------------------------------------- face detector
# FaceDetector: our own single-shot face + 5-landmark detector (trained in training/). Ideas from
# MTCNN (Zhang et al. 2016): three tasks per candidate (face / box / landmarks), PReLU, online
# hard-negative mining. Unlike MTCNN: one fully-convolutional network with a feature pyramid
# (strides 8/16/32) and anchors instead of an image pyramid + three cascaded networks.
DETECTOR_WEIGHTS = "face_detector.safetensors"
DETECTOR_WIDTH = 32
DETECT_THRESHOLD = 0.5        # face score above which a detection counts
ANCHOR_LEVELS = ((8, (16, 24, 32)), (16, (48, 64, 96)), (32, (128, 192, 256, 384)))   # stride, anchor sides
BOX_VAR = (0.1, 0.2)


def _conv_bn(cin, cout, k=3, s=1, groups=1):
    return nn.Sequential(nn.Conv2d(cin, cout, k, s, k // 2, groups=groups, bias=False), nn.BatchNorm2d(cout),
                         nn.PReLU(cout))


class _Block(nn.Module):
    """Residual block: (3x3 depthwise + 1x1 pointwise) twice; the stride is on the first depthwise."""

    def __init__(self, cin, cout, stride):
        super().__init__()
        self.body = nn.Sequential(_conv_bn(cin, cin, 3, stride, groups=cin), _conv_bn(cin, cout, 1),
                                  _conv_bn(cout, cout, 3, 1, groups=cout), _conv_bn(cout, cout, 1))
        self.skip = None if stride == 1 and cin == cout else nn.Sequential(
            nn.AvgPool2d(stride) if stride > 1 else nn.Identity(), nn.Conv2d(cin, cout, 1, bias=False), nn.BatchNorm2d(cout))

    def forward(self, x):
        return self.body(x) + (x if self.skip is None else self.skip(x))


class FaceDetector(nn.Module):
    """(N, 3, H, W) RGB 0..1 (H, W multiples of 32) -> per anchor: face logit, box deltas (4) and
    landmark offsets (10). `detect()` turns them into faces for one PIL image."""

    def __init__(self, width: int = DETECTOR_WIDTH):
        super().__init__()
        w = width
        self.stem = nn.Sequential(_conv_bn(3, w, 3, 2), _conv_bn(w, w, 3, 1, groups=w), _conv_bn(w, w, 1))   # /2
        self.s1 = nn.Sequential(_Block(w, 2 * w, 2), _Block(2 * w, 2 * w, 1))                                  # /4
        self.s2 = nn.Sequential(_Block(2 * w, 4 * w, 2), _Block(4 * w, 4 * w, 1), _Block(4 * w, 4 * w, 1))      # /8
        self.s3 = nn.Sequential(_Block(4 * w, 8 * w, 2), _Block(8 * w, 8 * w, 1), _Block(8 * w, 8 * w, 1))      # /16
        self.s4 = nn.Sequential(_Block(8 * w, 16 * w, 2), _Block(16 * w, 16 * w, 1))                            # /32
        f = 4 * w
        self.lat = nn.ModuleList([nn.Conv2d(c, f, 1) for c in (4 * w, 8 * w, 16 * w)])
        self.smooth = nn.ModuleList([_conv_bn(f, f) for _ in range(3)])
        self.context = nn.ModuleList([nn.Sequential(_conv_bn(f, f), _conv_bn(f, f, 3, 1, groups=f), _conv_bn(f, f, 1))
                                      for _ in range(3)])
        self.n_anchor = [len(sizes) for _, sizes in ANCHOR_LEVELS]
        self.cls = nn.ModuleList([nn.Conv2d(f, a, 1) for a in self.n_anchor])
        self.box = nn.ModuleList([nn.Conv2d(f, a * 4, 1) for a in self.n_anchor])
        self.ldm = nn.ModuleList([nn.Conv2d(f, a * 10, 1) for a in self.n_anchor])
        for c in self.cls:
            nn.init.constant_(c.bias, -4.595)       # log(0.01 / 0.99): start at p(face) = 0.01

    def forward(self, x):
        x = (x - 0.5) / 0.5
        c2 = self.s2(self.s1(self.stem(x)))
        c3 = self.s3(c2)
        c4 = self.s4(c3)
        p5 = self.lat[2](c4)
        p4 = self.lat[1](c3) + F.interpolate(p5, size=c3.shape[-2:], mode="nearest")
        p3 = self.lat[0](c2) + F.interpolate(p4, size=c2.shape[-2:], mode="nearest")
        outs = []
        for i, p in enumerate((p3, p4, p5)):
            p = self.smooth[i](p)
            p = p + self.context[i](p)
            n, _, h, w = p.shape
            a = self.n_anchor[i]
            outs.append((self.cls[i](p).permute(0, 2, 3, 1).reshape(n, h * w * a),
                         self.box[i](p).permute(0, 2, 3, 1).reshape(n, h * w * a, 4),
                         self.ldm[i](p).permute(0, 2, 3, 1).reshape(n, h * w * a, 10)))
        return tuple(torch.cat([o[k] for o in outs], 1) for k in range(3))

    @torch.inference_mode()
    def detect(self, images: list[Image.Image], score: float = DETECT_THRESHOLD, nms_iou: float = 0.4,
               max_side: int = 1024, min_side: int = 320) -> list[tuple[np.ndarray, np.ndarray, np.ndarray]]:
        """Same-size PIL images -> per image (boxes (K, 4), scores (K,), landmarks (K, 5, 2)) in
        image pixels, best first. The longer side is scaled into [min_side, max_side] first."""
        w0, h0 = images[0].size
        s = min(max_side / max(w0, h0), 1.0)
        if max(w0, h0) < min_side:          # small photos are upscaled (faces this small are hard to see)
            s = min_side / max(w0, h0)
        w, h = max(32, round(w0 * s / 32) * 32), max(32, round(h0 * s / 32) * 32)
        dev = next(self.parameters()).device
        x = torch.from_numpy(np.stack([np.asarray(im.resize((w, h), Image.BILINEAR), dtype=np.float32) for im in images]))
        logits, d_box, d_ldm = self(x.permute(0, 3, 1, 2).div(255).to(dev))
        anc = anchors(h, w, dev)
        sx, sy = w0 / w, h0 / h
        out = []
        for i in range(len(images)):
            p = logits[i].float().sigmoid()
            keep = p > score
            boxes, ldm = decode_detections(d_box[i][keep].float(), d_ldm[i][keep].float(), anc[keep])
            p = p[keep]
            k = torchvision.ops.nms(boxes, p, nms_iou)
            boxes = (boxes[k] * torch.tensor([sx, sy, sx, sy], device=dev)).cpu().numpy()
            ldm = (ldm[k].view(-1, 5, 2) * torch.tensor([sx, sy], device=dev)).cpu().numpy()
            out.append((boxes, p[k].cpu().numpy(), ldm))
        return out


_ANCHORS: dict = {}


def anchors(h: int, w: int, device) -> torch.Tensor:
    """(A, 4) anchors (cx, cy, side, side) for an h x w input, in the order of FaceDetector's outputs."""
    key = (h, w, str(device))
    if key not in _ANCHORS:
        out = []
        for stride, sizes in ANCHOR_LEVELS:
            fh, fw = -(-h // stride), -(-w // stride)
            ys, xs = torch.meshgrid(torch.arange(fh), torch.arange(fw), indexing="ij")
            cx = (xs.reshape(-1, 1).float() + 0.5) * stride
            cy = (ys.reshape(-1, 1).float() + 0.5) * stride
            side = torch.tensor(sizes, dtype=torch.float32).view(1, -1)
            out.append(torch.stack([cx.expand(-1, len(sizes)), cy.expand(-1, len(sizes)), side.expand(fh * fw, -1),
                                    side.expand(fh * fw, -1)], -1).reshape(-1, 4))
        _ANCHORS[key] = torch.cat(out).to(device)
    return _ANCHORS[key]


def decode_detections(d_box, d_ldm, anc):
    """Network deltas + anchors -> boxes (K, 4) x1 y1 x2 y2 and landmarks (K, 10)."""
    cx = anc[:, 0] + d_box[:, 0] * BOX_VAR[0] * anc[:, 2]
    cy = anc[:, 1] + d_box[:, 1] * BOX_VAR[0] * anc[:, 3]
    w = anc[:, 2] * torch.exp((d_box[:, 2] * BOX_VAR[1]).clamp(max=6))
    h = anc[:, 3] * torch.exp((d_box[:, 3] * BOX_VAR[1]).clamp(max=6))
    boxes = torch.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], 1)
    ldm = anc[:, None, :2] + d_ldm.view(-1, 5, 2) * BOX_VAR[0] * anc[:, None, 2:]
    return boxes, ldm.reshape(-1, 10)


def encode_detections(boxes, ldm, anc):
    """Inverse of decode_detections: training targets for anchors matched to faces."""
    cx, cy = (boxes[:, 0] + boxes[:, 2]) / 2, (boxes[:, 1] + boxes[:, 3]) / 2
    w, h = (boxes[:, 2] - boxes[:, 0]).clamp(min=1), (boxes[:, 3] - boxes[:, 1]).clamp(min=1)
    t_box = torch.stack([(cx - anc[:, 0]) / (BOX_VAR[0] * anc[:, 2]), (cy - anc[:, 1]) / (BOX_VAR[0] * anc[:, 3]),
                         torch.log(w / anc[:, 2]) / BOX_VAR[1], torch.log(h / anc[:, 3]) / BOX_VAR[1]], 1)
    t_ldm = (ldm.view(-1, 5, 2) - anc[:, None, :2]) / (BOX_VAR[0] * anc[:, None, 2:])
    return t_box, t_ldm.reshape(-1, 10)


def pick_face(boxes: np.ndarray, image_size, selection: str = "largest") -> int | None:
    """Which detection to use: the largest face, or (LFW) the most central large face."""
    if len(boxes) == 0:
        return None
    area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    if selection == "largest":
        return int(area.argmax())
    w, h = image_size
    offset = ((boxes[:, 0] + boxes[:, 2]) / 2 - w / 2) ** 2 + ((boxes[:, 1] + boxes[:, 3]) / 2 - h / 2) ** 2
    return int((area - 2.0 * offset).argmax())


def load_detector(model_dir: str | Path, device: str = "cpu") -> FaceDetector:
    det = FaceDetector()
    det.load_state_dict(load_file(_require(Path(model_dir) / DETECTOR_WEIGHTS)), strict=True)
    return det.to(device).eval()


# The standard 5-point face template (eyes, nose tip, mouth corners) of 112x112 face crops,
# scaled to `crop_size`. Every face is rotated/scaled/shifted onto it, so eyes and mouth always
# land on the same pixels.
FACE_TEMPLATE_112 = np.array([[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
                              [41.5493, 92.3655], [70.7299, 92.2041]])


def similarity_transform(src: np.ndarray, dst: np.ndarray) -> np.ndarray:
    """2x3 matrix of the rotation + uniform scale + shift that best maps `src` points onto `dst`
    (least squares, Umeyama 1991)."""
    mu_s, mu_d = src.mean(0), dst.mean(0)
    s, d = src - mu_s, dst - mu_d
    U, S, Vt = np.linalg.svd(d.T @ s / len(src))
    D = np.eye(2)
    if np.linalg.det(U) * np.linalg.det(Vt) < 0:
        D[1, 1] = -1
    R = U @ D @ Vt
    scale = (S * np.diag(D)).sum() / s.var(0).sum()
    return np.hstack([scale * R, (mu_d - scale * R @ mu_s)[:, None]])


def warp_face(img: Image.Image, landmarks: np.ndarray, crop_size: int) -> Image.Image:
    """Rotate/scale the photo so its 5 landmarks sit on the face template -> crop_size x crop_size."""
    M = similarity_transform(np.asarray(landmarks, dtype=np.float64), FACE_TEMPLATE_112 * crop_size / 112)
    inverse = np.linalg.inv(np.vstack([M, [0, 0, 1]]))[:2].ravel()     # PIL maps output -> input
    return img.transform((crop_size, crop_size), Image.AFFINE, tuple(inverse), resample=Image.BILINEAR)


def pil_to_tensor(img: Image.Image) -> torch.Tensor:
    return torch.from_numpy(np.asarray(img, dtype=np.float32).copy()).permute(2, 0, 1)


class FaceAligner:
    """Our FaceDetector (face_detector.safetensors) finds the face and its 5 landmarks; the face is
    then warped onto the template. image -> aligned (3, crop, crop) face tensor with values 0..255."""

    def __init__(self, model_dir: str | Path, device: str = "cpu", crop_size: int = 128,
                 selection: str = "largest", detector: FaceDetector | None = None):
        self.crop_size = crop_size
        self.selection = selection      # "largest" for photos; "center" for LFW (labelled by the centered person)
        self.detector = detector if detector is not None else load_detector(model_dir, device)

    def align_batch(self, images: list[Image.Image]) -> tuple[list[Image.Image], list[bool]]:
        """Same-size PIL images -> (aligned PIL faces, whether a face was found). No face -> center crop."""
        faces, found = [], []
        for img, (boxes, _, ldm) in zip(images, self.detector.detect(images)):
            k = pick_face(boxes, img.size, self.selection)
            if k is None:
                faces.append(center_crop_face(img, self.crop_size)); found.append(False)
            else:
                faces.append(warp_face(img, ldm[k], self.crop_size)); found.append(True)
        return faces, found

    def __call__(self, image) -> tuple[torch.Tensor, bool]:
        """Image -> (CPU face tensor (3, crop, crop) 0..255, whether a face was found)."""
        faces, found = self.align_batch([_limit_size(to_pil(image))])
        return pil_to_tensor(faces[0]), found[0]


class FaceEmbedder:
    """Aligner + the trained embedders: image -> (embedding, aligned face, face detected)."""

    def __init__(self, net: FaceRecognizer, aligner: FaceAligner, device: torch.device):
        self.net, self.align, self.device = net, aligner, device

    @torch.inference_mode()
    def embed_faces(self, faces: torch.Tensor) -> np.ndarray:
        """(N, 3, crop, crop) faces 0..255 -> (N, 1024) float32 L2-normalised embeddings."""
        return self.net.embed(faces.to(self.device)).float().cpu().numpy()

    def embed(self, image) -> tuple[np.ndarray, torch.Tensor, bool]:
        face, detected = self.align(image)
        return self.embed_faces(face.unsqueeze(0))[0], face, detected


# --------------------------------------------------------------------------- inference API
class Predictor:
    """Loads detector, embedder and identity head once; identifies, verifies and embeds faces."""

    def __init__(self, model_dir: str | Path, device: str = "cpu"):
        model_dir = Path(model_dir)
        _require(model_dir / WEIGHTS)
        self.device = torch.device(device)
        self.net = FaceRecognizer.from_pretrained(model_dir, map_location="cpu", strict=True)
        self.net.to(self.device).eval()
        self.class_names = self.net.class_names
        self.verify_threshold = self.net.verify_threshold
        aligner = FaceAligner(model_dir, str(self.device), self.net.crop_size)
        self.embedder = FaceEmbedder(self.net, aligner, self.device)

    @torch.inference_mode()
    def predict_embeddings(self, embeddings: np.ndarray) -> np.ndarray:
        """(N, 1024) embeddings -> (N, n_classes) softmax probabilities (used by evaluation)."""
        x = torch.as_tensor(np.asarray(embeddings, dtype=np.float32), device=self.device)
        return self.net.classify(x).softmax(dim=1).float().cpu().numpy()

    def _probs_dict(self, embedding: np.ndarray) -> dict[str, float]:
        probs = self.predict_embeddings(embedding[None])[0]
        order = np.argsort(-probs)
        return {self.class_names[i]: float(probs[i]) for i in order}

    def predict(self, image) -> dict[str, float]:
        """Face photo -> {person: probability} for all known people, most likely first (sums to 1)."""
        embedding, _, _ = self.embedder.embed(image)
        return self._probs_dict(embedding)

    def analyze(self, image) -> dict:
        """Everything the Space shows: probabilities, the aligned crop, whether a face was detected
        and the embedding (numpy)."""
        embedding, face, detected = self.embedder.embed(image)
        return {"probs": self._probs_dict(embedding), "face": face_to_pil(face),
                "face_detected": detected, "embedding": embedding}

    def embed(self, image) -> list[float]:
        """Face photo -> 1024-d L2-normalised embedding (plain floats)."""
        return self.embedder.embed(image)[0].astype(float).tolist()

    def verify(self, image_a, image_b, threshold: float | None = None) -> dict:
        """Are two photos the same person? Cosine distance of their embeddings vs `threshold`."""
        threshold = self.verify_threshold if threshold is None else float(threshold)
        emb_a, _, found_a = self.embedder.embed(image_a)
        emb_b, _, found_b = self.embedder.embed(image_b)
        distance = cosine_distance(emb_a, emb_b)
        return {"same_person": bool(distance < threshold), "distance": round(distance, 6),
                "threshold": threshold, "faces_detected": [found_a, found_b]}


def load(model_dir: str | Path, device: str = "cpu") -> Predictor:
    return Predictor(model_dir, device)