engine.py

"""
app.py — Core computer vision engine for the Passport Photo Maker.

Pipeline: upload -> downscale guard -> BiRefNet-lite segmentation ->
YuNet face detection -> geometric crop to spec -> background composite ->
300 DPI canvas render -> print sheet tiling.

Hard constraint: must run stably on a free HF CPU Basic Space (2 vCPU,
16GB RAM, no swap headroom to waste). Every design choice below exists
because of that constraint — read the comments before "optimizing" them
away.
"""

from __future__ import annotations

import gc
import io
import logging
import os
import threading
from dataclasses import dataclass
from typing import Optional

import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

try:
    import spaces  # HF ZeroGPU runtime — only importable/meaningful on Spaces
    _ZEROGPU = True
except ImportError:
    _ZEROGPU = False

    class _NoOpSpaces:
        """Local/non-Spaces fallback so `@spaces.GPU` is a harmless no-op
        decorator when running outside HF ZeroGPU (e.g. local dev, or a
        future move back to CPU Basic hardware)."""

        @staticmethod
        def GPU(fn=None, **kwargs):
            if fn is None:
                return lambda f: f
            return fn

    spaces = _NoOpSpaces()

DEVICE = "cuda" if (_ZEROGPU and torch.cuda.is_available()) else "cpu"

# ---------------------------------------------------------------------------
# Global CPU tuning — must run before any heavy torch op.
# HF CPU Basic = 2 vCPUs. PyTorch defaults to detecting all cores, which on
# a shared/throttled container causes thread oversubscription (threads
# fighting each other for the same 2 cores -> slower AND memory-spikier
# because more intermediate buffers are alive concurrently). Pin it down.
# Only relevant on CPU hardware — on ZeroGPU the model runs on the A10G,
# so CPU thread starvation isn't the bottleneck and this is skipped.
# ---------------------------------------------------------------------------
if DEVICE == "cpu":
    torch.set_num_threads(2)
    torch.set_num_interop_threads(1)

MODEL_ID = "ZhengPeng7/BiRefNet_lite"
MODEL_REVISION = "7838f1c3472f827cd8ce13ab5ccc2ce48077360f"  # pinned commit —
# this model loads with trust_remote_code=True, meaning the model
# author's Python code (birefnet.py) executes directly in this process.
# Without a pinned revision, a future push to the model repo's "main"
# branch would auto-execute here on next cold start with no review.
# Bump this hash deliberately (and re-check birefnet.py) if the model
# needs updating — never leave this unpinned with trust_remote_code=True.
SEG_INPUT_SIZE = 1024  # BiRefNet's trained resolution. Do NOT drop to 512 —
# that is a training-resolution mismatch, not a speed optimization; it
# degrades mask quality (soft/wrong edges) for a marginal CPU saving that
# downscaling the *source* image already captures. Speed comes from the
# 1200px container cap below, not from starving the segmentation model.

MAX_CONTAINER_PX = 1200  # hard cap on the longest side of any uploaded
# image before it touches any model. This is the actual OOM guard — a
# 12MP phone photo run through a segmentation model on 16GB shared RAM
# is what crashes Spaces, not the model itself.

DPI = 300


# ---------------------------------------------------------------------------
# Model singleton — loaded once per worker process, never per-request.
# Reloading a ~200M param model on every click is the #1 cause of slow /
# OOM-prone HF Spaces. Thread lock because Gradio's queue can dispatch
# concurrent requests (see ui.py concurrency_limit=2) onto the same process.
# ---------------------------------------------------------------------------
_model_lock = threading.Lock()
_model: Optional[torch.nn.Module] = None

_seg_transform = transforms.Compose(
    [
        transforms.Resize((SEG_INPUT_SIZE, SEG_INPUT_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)


def get_model() -> torch.nn.Module:
    """Lazy-load and cache the BiRefNet-lite segmentation model.

    Deliberately eager mode (no torch.jit.trace). BiRefNet-lite's decoder
    has data-dependent branching across its multi-stage refinement heads;
    tracing captures ONE execution path against the dummy input and will
    silently produce wrong output shapes/masks on inputs that take a
    different path. That is a correctness bug, not a performance one —
    not worth the trade for a speedup eager mode mostly already gets from
    torch.no_grad() + thread pinning + input-size discipline.
    """
    global _model
    if _model is None:
        with _model_lock:
            if _model is None:  # double-checked locking
                m = AutoModelForImageSegmentation.from_pretrained(
                    MODEL_ID, revision=MODEL_REVISION, trust_remote_code=True
                )
                m.eval()
                # Load on CPU regardless of target device. ZeroGPU only
                # attaches a real CUDA device to the process INSIDE a
                # @spaces.GPU-decorated call — touching .to("cuda") here,
                # at load/import time, happens outside that context and
                # will fail (no GPU attached yet). The move to CUDA (when
                # DEVICE == "cuda") happens per-call in segment_alpha(),
                # which is itself the decorated function.
                m.to("cpu")
                # Disable gradient tracking at the parameter level too —
                # belt-and-suspenders on top of the no_grad() context used
                # at inference time; keeps autograd graph bookkeeping off
                # entirely for this process's lifetime.
                for p in m.parameters():
                    p.requires_grad_(False)
                _model = m
    return _model


def warm_up() -> None:
    """Run one dummy inference at import time so the FIRST real user
    request isn't the one eating model-load + cold-kernel latency.
    Call this once from app.py at Space boot, not per-request.

    Deliberately CPU-only regardless of DEVICE: on ZeroGPU, no CUDA
    device is attached to the process at boot time — it's only attached
    inside a live @spaces.GPU-decorated call during an actual user
    request. So this warms up model-loading + weight-download only; the
    first real request still pays GPU-attach latency on ZeroGPU (a few
    seconds), which is a ZeroGPU platform cost, not something app code
    can avoid.
    """
    model = get_model()
    dummy = torch.zeros(1, 3, SEG_INPUT_SIZE, SEG_INPUT_SIZE)
    with torch.no_grad():
        _ = model(dummy)  # runs on CPU — model was loaded via .to("cpu")
    del dummy
    gc.collect()
    _ensure_yunet_model()  # pre-download face detector weights too

    # Pre-download + warm the pose model too (outfit overlay feature).
    # Wrapped in try/except: outfit overlay is an optional enhancement,
    # so a warm-up failure here (e.g. onnxruntime missing, HF Hub hiccup)
    # should not crash Space boot — apply_outfit()'s own error handling
    # already degrades gracefully per-request if the pose model is
    # unavailable.
    try:
        session = _get_pose_session()
        dummy_pose = np.zeros((1, _MOVENET_INPUT_SIZE, _MOVENET_INPUT_SIZE, 3), dtype=np.int32)
        session.run(None, {session.get_inputs()[0].name: dummy_pose})
        del dummy_pose
        gc.collect()
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Standards matrix
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class PhotoSpec:
    label: str
    width_mm: float
    height_mm: float
    bg_hex: str  # default/expected background color for this spec
    head_ratio: float  # target head-height as a fraction of photo height
    # (standard passport convention: head occupies ~70-80% of frame height
    # measured chin-to-crown; we default to 0.70 "breathing room" per spec)


def _spec(label: str, w: float, h: float, bg: str, ratio: float) -> PhotoSpec:
    return PhotoSpec(label, w, h, bg, ratio)


# Head-ratio note: this is head-height / frame-height, used against
# YuNet's expanded crown-to-chin box in crop_to_spec(). Country columns
# below give a chin-to-crown mm range within the frame; ratio here is
# that range's midpoint divided by the frame height (mm), which is the
# same math the original 4 specs used at 0.70 — not a fresh assumption.
STANDARDS: dict[str, PhotoSpec] = {
    # --- Pakistan (kept from original) ---
    "Pakistan Passport (Blue, 38x51mm)": _spec(
        "Pakistan Passport (Blue, 38x51mm)", 38.0, 51.0, "#1E3A8A", 0.70
    ),
    "Pakistan CNIC / NADRA (White, 38x51mm)": _spec(
        "Pakistan CNIC / NADRA (White, 38x51mm)", 38.0, 51.0, "#FFFFFF", 0.70
    ),
    # --- 51x51mm / 2x2in family (US, India, Philippines etc share size, NOT specs) ---
    "USA Passport / Visa / DS-160 (White, 2x2in / 51x51mm)": _spec(
        "USA Passport / Visa / DS-160 (White, 2x2in / 51x51mm)", 51.0, 51.0, "#FFFFFF", 0.62
    ),
    "India Passport / OCI (White, 51x51mm)": _spec(
        "India Passport / OCI (White, 51x51mm)", 51.0, 51.0, "#FFFFFF", 0.65
    ),
    "India PAN Card (White, 25x35mm)": _spec(
        "India PAN Card (White, 25x35mm)", 25.0, 35.0, "#FFFFFF", 0.65
    ),
    "Philippines Passport (White, 2x2in / 51x51mm)": _spec(
        "Philippines Passport (White, 2x2in / 51x51mm)", 51.0, 51.0, "#FFFFFF", 0.62
    ),
    "Brazil Visa (White, 51x51mm)": _spec(
        "Brazil Visa (White, 51x51mm)", 51.0, 51.0, "#FFFFFF", 0.62
    ),
    # --- 35x45mm family (UK, Schengen, most of world) ---
    "UK Passport / Visa (Light Grey, 35x45mm)": _spec(
        "UK Passport / Visa (Light Grey, 35x45mm)", 35.0, 45.0, "#E8E8E8", 0.71
    ),
    "UKVI Visa (White, 35x45mm)": _spec(
        "UKVI Visa (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Schengen Visa — EU/France/Germany/Italy/Spain (Light Grey, 35x45mm)": _spec(
        "Schengen Visa — EU/France/Germany/Italy/Spain (Light Grey, 35x45mm)", 35.0, 45.0, "#F0F0F0", 0.71
    ),
    "Ireland Passport (Light Grey, 35x45mm)": _spec(
        "Ireland Passport (Light Grey, 35x45mm)", 35.0, 45.0, "#E8E8E8", 0.71
    ),
    "Germany Passport / Biometric ID (Light Grey, 35x45mm)": _spec(
        "Germany Passport / Biometric ID (Light Grey, 35x45mm)", 35.0, 45.0, "#E8E8E8", 0.71
    ),
    "India Visa / Most Documents (White, 35x45mm)": _spec(
        "India Visa / Most Documents (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Australia Passport (White, 35x45mm)": _spec(
        "Australia Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "New Zealand Passport (White, 35x45mm)": _spec(
        "New Zealand Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Japan Passport / Visa (White, 35x45mm)": _spec(
        "Japan Passport / Visa (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "South Korea Passport (White, 35x45mm)": _spec(
        "South Korea Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Singapore Passport (White, 35x45mm)": _spec(
        "Singapore Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Russia Passport / Visa (White, 35x45mm)": _spec(
        "Russia Passport / Visa (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Bangladesh Passport (White, 35x45mm)": _spec(
        "Bangladesh Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Nigeria Passport (White, 35x45mm)": _spec(
        "Nigeria Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    # --- Unique-size outliers ---
    "Canada Passport / Visa (White, 50x70mm)": _spec(
        "Canada Passport / Visa (White, 50x70mm)", 50.0, 70.0, "#FFFFFF", 0.48
    ),
    "Brazil Passport (White, 50x70mm)": _spec(
        "Brazil Passport (White, 50x70mm)", 50.0, 70.0, "#FFFFFF", 0.48
    ),
    "China Passport (White, 33x48mm)": _spec(
        "China Passport (White, 33x48mm)", 33.0, 48.0, "#FFFFFF", 0.65
    ),
    "China Visa (Light Blue, 33x48mm)": _spec(
        "China Visa (Light Blue, 33x48mm)", 33.0, 48.0, "#C6E2F5", 0.65
    ),
    "UAE Visa (White, 43x55mm)": _spec(
        "UAE Visa (White, 43x55mm)", 43.0, 55.0, "#FFFFFF", 0.68
    ),
    "Saudi Arabia Visa (White, 40x60mm)": _spec(
        "Saudi Arabia Visa (White, 40x60mm)", 40.0, 60.0, "#FFFFFF", 0.65
    ),
    "Malaysia Passport / Visa (Blue, 35x50mm)": _spec(
        "Malaysia Passport / Visa (Blue, 35x50mm)", 35.0, 50.0, "#1E3A8A", 0.65
    ),
    "Spain National ID / DNI (White, 26x32mm)": _spec(
        "Spain National ID / DNI (White, 26x32mm)", 26.0, 32.0, "#FFFFFF", 0.71
    ),
    "Mexico Passport (White, 35x45mm)": _spec(
        "Mexico Passport (White, 35x45mm)", 35.0, 45.0, "#FFFFFF", 0.71
    ),
    "Vietnam Passport / Visa (White, 40x60mm)": _spec(
        "Vietnam Passport / Visa (White, 40x60mm)", 40.0, 60.0, "#FFFFFF", 0.65
    ),
    "Turkey Passport / Visa (White, 50x60mm)": _spec(
        "Turkey Passport / Visa (White, 50x60mm)", 50.0, 60.0, "#FFFFFF", 0.65
    ),
}


def mm_to_px(mm: float, dpi: int = DPI) -> int:
    """Exact mm -> px conversion at a fixed DPI. 1 inch = 25.4mm."""
    return round((mm / 25.4) * dpi)


# ---------------------------------------------------------------------------
# Stage 1 — ingest guard
# ---------------------------------------------------------------------------
def load_and_bound(image: Image.Image, max_px: int = MAX_CONTAINER_PX) -> Image.Image:
    """Downscale any oversized upload to a max_px bounding box.

    Uses LANCZOS, not NEAREST. NEAREST is cheaper but produces aliased,
    jagged edges on the downscaled source — exactly the wrong input to
    feed a segmentation model (it amplifies edge noise the model then has
    to guess through) and the wrong input for a face crop that will be
    printed. LANCZOS costs low-single-digit milliseconds extra at this
    resolution; that is not where your CPU budget is being spent. Upscale
    guard uses BICUBIC for the same reason (never NEAREST for anything
    that ends up in a printed deliverable).
    """
    image = image.convert("RGB")
    w, h = image.size
    if max(w, h) <= max_px:
        return image
    scale = max_px / max(w, h)
    new_size = (max(1, round(w * scale)), max(1, round(h * scale)))
    return image.resize(new_size, Image.Resampling.LANCZOS)


# ---------------------------------------------------------------------------
# Stage 2 — background segmentation (BiRefNet-lite)
# ---------------------------------------------------------------------------
@spaces.GPU(duration=20)  # ZeroGPU: request up to 20s of A10G time per call.
# Anything touching CUDA (model.to("cuda"), tensor ops on a GPU tensor)
# must happen INSIDE this decorated function — ZeroGPU only attaches a
# real device to the process for the duration of a decorated call, then
# revokes it. On plain CPU Basic hardware (DEVICE == "cpu") this
# decorator is a harmless no-op (see the _NoOpSpaces fallback above).
def segment_alpha(image: Image.Image) -> Image.Image:
    """Return an RGBA image with an accurate alpha matte cut around the
    subject. Strict memory hygiene: every intermediate tensor is deleted
    and gc.collect() is called immediately after the numpy/PIL handoff,
    because this is the single most memory-heavy step in the pipeline and
    the one most likely to trigger an OOM under concurrent requests (or,
    on ZeroGPU, VRAM pressure on the shared A10G pool).
    """
    model = get_model()
    if DEVICE == "cuda":
        model = model.to("cuda")
    orig_size = image.size  # (W, H)

    inp = _seg_transform(image).unsqueeze(0)  # (1,3,1024,1024)
    if DEVICE == "cuda":
        inp = inp.to("cuda")

    with torch.no_grad():
        preds = model(inp)
        # BiRefNet returns a list of side-outputs across decoder stages;
        # the final, highest-resolution prediction is last.
        pred = preds[-1] if isinstance(preds, (list, tuple)) else preds
        pred = pred.sigmoid().squeeze()  # (1024,1024) in [0,1]

    mask_np = (pred.cpu().numpy() * 255).astype(np.uint8)

    # --- strict memory wipe of everything torch-side, immediately ---
    del inp, preds, pred
    if DEVICE == "cuda":
        # Move model back off GPU memory between calls — ZeroGPU revokes
        # device access after the decorated function returns anyway, but
        # explicit cache clearing avoids leaving stale allocations that
        # count against the shared pool while this worker is between
        # requests.
        model.to("cpu")
        torch.cuda.empty_cache()
    gc.collect()

    mask_img = Image.fromarray(mask_np).resize(orig_size, Image.Resampling.LANCZOS)
    del mask_np
    gc.collect()

    result = image.convert("RGBA")
    result.putalpha(mask_img)
    del mask_img
    gc.collect()

    return result


# ---------------------------------------------------------------------------
# Stage 3 — face localization (YuNet, not Haar)
# ---------------------------------------------------------------------------
# Why YuNet instead of the Haar cascade the original spec asked for:
# Haar (haarcascade_frontalface_default.xml) is template-matching-era CV.
# It fails hard on: tilted/rotated heads, side lighting, glasses glare,
# partial occlusion, non-frontal yaw beyond ~15deg — all things that show
# up constantly in real phone-camera passport-photo uploads. A tool meant
# to compete with Cutout.pro cannot ship a face detector that fails on a
# meaningful slice of real uploads.
#
# YuNet (cv2.FaceDetectorYN) is still "classic OpenCV" — it ships inside
# opencv-contrib as a bundled ~230KB ONNX model run through OpenCV's own
# lightweight DNN backend. It is NOT a heavy framework like mediapipe or
# a face-recognition deep pipeline: no separate runtime, no extra Python
# package, single-digit-millisecond CPU inference, and materially better
# recall/robustness on tilt, small faces, and partial occlusion. This is
# the correct trade for "zero-RAM-framework-weight but actually works."
_YUNET_MODEL_URL = (
    "https://github.com/opencv/opencv_zoo/raw/main/models/"
    "face_detection_yunet/face_detection_yunet_2023mar.onnx"
)
_YUNET_MODEL_PATH = os.path.join(
    os.path.dirname(__file__), "models", "face_detection_yunet_2023mar.onnx"
)

_face_detector_lock = threading.Lock()
_face_detector: Optional[cv2.FaceDetectorYN] = None


def _ensure_yunet_model() -> None:
    """Download the YuNet ONNX weights (~230KB, BSD-3, OpenCV Zoo) on
    first run if not already present. Plain Gradio SDK Spaces have no
    Docker build step to pre-cache this in, so it downloads once at
    process start instead — negligible one-time cost (~1s on Spaces'
    network), then reused for the container's lifetime.
    """
    if os.path.exists(_YUNET_MODEL_PATH):
        return
    os.makedirs(os.path.dirname(_YUNET_MODEL_PATH), exist_ok=True)
    import urllib.request

    urllib.request.urlretrieve(_YUNET_MODEL_URL, _YUNET_MODEL_PATH)


def _get_face_detector(input_size: tuple[int, int]) -> cv2.FaceDetectorYN:
    global _face_detector
    with _face_detector_lock:
        if _face_detector is None:
            _ensure_yunet_model()
            _face_detector = cv2.FaceDetectorYN.create(
                _YUNET_MODEL_PATH,
                "",
                input_size,
                score_threshold=0.7,
                nms_threshold=0.3,
                top_k=10,
            )
        else:
            _face_detector.setInputSize(input_size)
    return _face_detector


@dataclass(frozen=True)
class FaceBox:
    x: int
    y: int
    w: int
    h: int
    # Optional: raw YuNet landmarks (right_eye, left_eye, nose, right_mouth,
    # left_mouth) as (x,y) tuples, and detection confidence 0-1. Both
    # default to None/0.0 so every existing call site that only unpacks
    # x/y/w/h keeps working unchanged — only check_compliance() reads
    # these two fields.
    landmarks: Optional[tuple[tuple[float, float], ...]] = None
    score: float = 0.0

    @property
    def cx(self) -> float:
        return self.x + self.w / 2

    @property
    def cy(self) -> float:
        return self.y + self.h / 2


def detect_main_face(image_rgb: Image.Image) -> FaceBox:
    """Detect the primary (largest-area) face in the image via YuNet.

    Raises ValueError if no face is found — callers must surface this to
    the user rather than silently center-cropping a photo with no
    detected face, which would produce a passport photo that fails
    real-world verification.
    """
    np_img = np.array(image_rgb.convert("RGB"))[:, :, ::-1]  # RGB -> BGR
    h, w = np_img.shape[:2]

    detector = _get_face_detector((w, h))
    _, faces = detector.detect(np_img)

    del np_img
    gc.collect()

    if faces is None or len(faces) == 0:
        raise ValueError(
            "No face detected in the uploaded photo. Please upload a clear, "
            "front-facing photo with good lighting."
        )

    # faces: Nx15 array, cols 0:4 = x,y,w,h, cols 4:14 = 5 landmark (x,y)
    # pairs (right eye, left eye, nose tip, right mouth corner, left mouth
    # corner), col 14 = detection confidence score.
    areas = faces[:, 2] * faces[:, 3]
    best = faces[int(np.argmax(areas))]
    x, y, fw, fh = best[0:4]
    landmarks = tuple((float(best[4 + 2 * i]), float(best[5 + 2 * i])) for i in range(5))
    score = float(best[14])

    # Clamp to image bounds — YuNet can return slightly negative/over-edge
    # boxes near frame borders.
    x = max(0, int(round(x)))
    y = max(0, int(round(y)))
    fw = min(int(round(fw)), w - x)
    fh = min(int(round(fh)), h - y)

    return FaceBox(x, y, fw, fh, landmarks=landmarks, score=score)


# ---------------------------------------------------------------------------
# Stage 3.5 — auto-straighten (head AND body, via whole-frame rotation)
# ---------------------------------------------------------------------------
# Why whole-frame rotation instead of a head-only transform: a photo is one
# rigid rectangle of pixels. A person's head and shoulders/torso are also
# rigidly connected in nearly all passport-photo poses (nobody photographs
# themselves with their neck bent sideways relative to their own shoulders
# — that would look broken). So the single eye-line angle that describes
# "how tilted is the head" is, for a straight-postured subject, the SAME
# angle that describes "how tilted is the whole body" in the frame.
# Rotating the entire bounded image by the negative of that angle around
# its center levels head, shoulders, and torso together in one operation —
# which is what "straighten the person" means for a rigid photograph.
# This is NOT a per-limb pose-warp (that would need a body pose model —
# see engine.py history/skill notes for why that's deferred) — it is a
# whole-image rotation, the correct and only physically consistent
# operation for a single ridid-body subject in a 2D photo.
MAX_AUTO_STRAIGHTEN_DEG = 25.0  # safety ceiling — a detected "tilt" beyond
# this is more likely a face-detector landmark error than a real head tilt;
# rotating that far would crop away too much of the subject after the
# corners fall outside the frame. Beyond this, skip rotation and let the
# existing tilt compliance-check warn the user instead.


def compute_tilt_angle(face: FaceBox) -> float:
    """Eye-line angle in degrees from horizontal, positive = subject's
    right eye is higher than left eye in image coordinates. Returns 0.0
    if landmarks are unavailable (never fails the pipeline over this).
    """
    if not face.landmarks or len(face.landmarks) < 2:
        return 0.0
    import math

    (rx, ry), (lx, ly) = face.landmarks[0], face.landmarks[1]
    return math.degrees(math.atan2(ly - ry, lx - rx))


def straighten_image(image: Image.Image, angle_deg: float) -> Image.Image:
    """Rotate the whole frame to level the eye-line to horizontal. Expands
    the canvas (expand=True) so no corner content is clipped, then the
    caller re-runs face detection on the rotated result — rotating the
    old FaceBox coordinates analytically is more error-prone than just
    re-detecting on the already-cheap YuNet pass.

    Always returns RGB (never RGBA) — this runs on `bounded`, which is
    RGB going into segment_alpha()'s Normalize transform; that transform
    is hard-coded to 3 channels (see _seg_transform) and raises a shape
    RuntimeError on 4-channel input. The rotation itself is done in RGBA
    internally so PIL can fill the corners exposed by expand=True with a
    real transparent value instead of smearing edge pixels (BICUBIC
    without an alpha channel would otherwise blend rotated content
    against whatever garbage sits outside the original frame) — that
    RGBA intermediate is then flattened onto a neutral gray backing
    before return, so the corners become plain pixels the segmentation
    model can process like any other background, not transparency it has
    never seen and has no defined behavior for.

    BICUBIC resample, matching every other resize/rotate op in this file
    that ends up in a printed deliverable (see load_and_bound's docstring
    for why NEAREST is never used here).
    """
    if abs(angle_deg) < 0.5:
        return image  # sub-half-degree tilt isn't worth a resample pass
    angle_deg = max(-MAX_AUTO_STRAIGHTEN_DEG, min(MAX_AUTO_STRAIGHTEN_DEG, angle_deg))
    rgba = image.convert("RGBA")
    rotated = rgba.rotate(
        angle_deg,  # NOT negated — verified numerically: compute_tilt_angle()
        # returns math.atan2(ly-ry, lx-rx) in image (y-down) coordinates,
        # and PIL's Image.rotate(theta) rotates the image content by theta
        # measured in that same y-down convention, so passing the raw
        # measured angle (not its negative) is what levels the eye-line to
        # horizontal. A synthetic two-point test (eyes at a known 16.7°
        # tilt) confirmed rotate(+angle) drives the post-rotation eye-line
        # to ~0.0°, while rotate(-angle) doubles the tilt to ~33°. Do not
        # "simplify" this back to -angle_deg without re-running that check.
        resample=Image.Resampling.BICUBIC,
        expand=True,
        fillcolor=(0, 0, 0, 0),
    )

    # Flatten onto a neutral 50%-gray RGB canvas — not white/black, so the
    # exposed-corner triangles don't accidentally read as a "clean white
    # background" region to segment_alpha's foreground/background
    # separation (a pure white corner touching a light shirt could bias
    # the matte). Gray is a safe, low-signal fill any segmentation model
    # treats as unremarkable background.
    backing = Image.new("RGB", rotated.size, (128, 128, 128))
    backing.paste(rotated, mask=rotated.split()[3])  # alpha channel as mask
    return backing


# ---------------------------------------------------------------------------
# Stage 4 — geometric crop to passport spec
# ---------------------------------------------------------------------------
def compute_crop_box(
    img_w: int,
    img_h: int,
    face: FaceBox,
    spec: PhotoSpec,
    zoom: float = 1.0,
    x_offset: float = 0.0,
    y_offset: float = 0.0,
) -> tuple[float, float, float, float]:
    """Pure geometry: compute the (left, top, right, bottom) crop box
    crop_to_spec() will use, WITHOUT actually cropping/padding/resizing
    anything. Extracted as its own function so apply_outfit() can know
    exactly what region of `bounded`/`matted` will survive into the final
    photo BEFORE compositing a garment — critical because a passport
    photo's crop window only shows a small sliver below the chin, and a
    garment must be scaled to fit that sliver, not to shoulder width (see
    apply_outfit()'s docstring for why shoulder-width scaling produced an
    oversized, mostly-cropped-away garment in practice).
    """
    head_top = face.y - face.h * 0.55
    head_bottom = face.y + face.h * 1.35
    head_height = head_bottom - head_top
    head_cx = face.cx

    target_aspect = spec.width_mm / spec.height_mm
    crop_h = head_height / spec.head_ratio
    crop_w = crop_h * target_aspect

    head_cy = (head_top + head_bottom) / 2
    crop_top = head_cy - crop_h * 0.45
    crop_left = head_cx - crop_w / 2
    crop_bottom = crop_top + crop_h
    crop_right = crop_left + crop_w

    zoom = max(0.5, min(2.0, zoom))
    if zoom != 1.0 or x_offset or y_offset:
        box_cx = (crop_left + crop_right) / 2
        box_cy = (crop_top + crop_bottom) / 2
        new_w = crop_w / zoom
        new_h = crop_h / zoom
        box_cx += x_offset * new_w
        box_cy += y_offset * new_h
        crop_left = box_cx - new_w / 2
        crop_right = box_cx + new_w / 2
        crop_top = box_cy - new_h / 2
        crop_bottom = box_cy + new_h / 2

    return crop_left, crop_top, crop_right, crop_bottom


def crop_to_spec(
    image_rgba: Image.Image,
    face: FaceBox,
    spec: PhotoSpec,
    zoom: float = 1.0,
    x_offset: float = 0.0,
    y_offset: float = 0.0,
) -> Image.Image:
    """Crop/pad the segmented image so the face sits centered with correct
    vertical breathing room, then resize to the spec's exact 300 DPI pixel
    dimensions.

    zoom / x_offset / y_offset: manual override on top of the auto-computed
    crop box, driven by the UI's Adjust Crop sliders. zoom>1 tightens the
    box (zooms in), zoom<1 loosens it (zooms out, more headroom); offsets
    are fractions of crop_w/crop_h, so 0.1 shifts the box by 10% of its own
    size — this keeps the sliders' effect resolution-independent regardless
    of source photo size. Applied to the auto box, never replacing the
    head-ratio math, so a user who touches nothing gets identical output
    to before this param existed.

    Passport convention approximated here: head height (crown-to-chin,
    approximated from the face detector's bounding box height with a
    standard expansion factor since YuNet's box is eyes/nose/mouth-tight,
    not crown-to-chin) should occupy roughly `spec.head_ratio` of the
    final photo height, with the face vertically centered slightly above
    frame-center to leave correct shoulder/headroom balance.
    """
    img_w, img_h = image_rgba.size

    crop_left, crop_top, crop_right, crop_bottom = compute_crop_box(
        img_w, img_h, face, spec, zoom, x_offset, y_offset
    )

    # If the ideal crop extends beyond the source image, pad with the
    # spec's background color rather than shrinking the crop (which would
    # violate the head-ratio requirement). This keeps composition correct
    # even for tightly-framed source photos.
    pad_left = max(0, -crop_left)
    pad_top = max(0, -crop_top)
    pad_right = max(0, crop_right - img_w)
    pad_bottom = max(0, crop_bottom - img_h)

    if pad_left or pad_top or pad_right or pad_bottom:
        new_w = img_w + int(np.ceil(pad_left + pad_right))
        new_h = img_h + int(np.ceil(pad_top + pad_bottom))
        padded = Image.new("RGBA", (new_w, new_h), (0, 0, 0, 0))
        px, py = int(round(pad_left)), int(round(pad_top))
        padded.paste(image_rgba, (px, py))
        image_rgba = padded
        crop_left += px
        crop_top += py
        crop_right += px
        crop_bottom += py
        img_w, img_h = new_w, new_h

    crop_box = (
        int(round(crop_left)),
        int(round(crop_top)),
        int(round(crop_right)),
        int(round(crop_bottom)),
    )
    cropped = image_rgba.crop(crop_box)

    target_px = (mm_to_px(spec.width_mm), mm_to_px(spec.height_mm))
    # BICUBIC for the final resize — this is the deliverable pixel grid
    # at exact 300 DPI dimensions, quality matters more than the
    # microseconds NEAREST would save here.
    result = cropped.resize(target_px, Image.Resampling.BICUBIC)
    del cropped
    gc.collect()
    return result


# ---------------------------------------------------------------------------
# Stage 5 — background compositing
# ---------------------------------------------------------------------------
def composite_background(image_rgba: Image.Image, hex_color: str) -> Image.Image:
    """Flatten the alpha-matted subject onto a solid background color.
    Output is RGB (no alpha) since passport photo deliverables are
    printed/uploaded as flat JPEG/PNG without transparency.
    """
    hex_color = hex_color.lstrip("#")
    if len(hex_color) != 6:
        raise ValueError(f"Invalid hex color: #{hex_color}")
    rgb = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))

    bg = Image.new("RGBA", image_rgba.size, rgb + (255,))
    flattened = Image.alpha_composite(bg, image_rgba).convert("RGB")
    del bg
    gc.collect()
    return flattened


# ---------------------------------------------------------------------------
# Stage 6 — print sheet layout engine
# ---------------------------------------------------------------------------
PAPER_SIZES_MM = {
    "4x6 inch": (101.6, 152.4),
    "A4": (210.0, 297.0),
}

SHEET_MARGIN_MM = 3.0  # outer sheet margin
PHOTO_GUTTER_MM = 2.0  # spacing between tiled photos
BORDER_HEX = "#B0B0B0"  # subtle gray divider border
BORDER_WIDTH_PX = 1  # exact 1px width, independent of DPI scale, per spec


def build_print_sheet(photo: Image.Image, spec: PhotoSpec, paper: str) -> Image.Image:
    """Tile a single passport photo across a print sheet at exact pixel
    ratios (no stretching/distortion — every tile is a 1:1 pixel copy of
    the source photo, laid out on a grid computed from real mm math).
    """
    if paper not in PAPER_SIZES_MM:
        raise ValueError(f"Unknown paper size: {paper}")

    paper_w_mm, paper_h_mm = PAPER_SIZES_MM[paper]
    sheet_w_px = mm_to_px(paper_w_mm)
    sheet_h_px = mm_to_px(paper_h_mm)

    photo_w_px, photo_h_px = photo.size  # already exact 300 DPI spec size
    gutter_px = mm_to_px(PHOTO_GUTTER_MM)
    margin_px = mm_to_px(SHEET_MARGIN_MM)

    usable_w = sheet_w_px - 2 * margin_px
    usable_h = sheet_h_px - 2 * margin_px

    cols = max(1, (usable_w + gutter_px) // (photo_w_px + gutter_px))
    rows = max(1, (usable_h + gutter_px) // (photo_h_px + gutter_px))

    if cols == 0 or rows == 0:
        raise ValueError(
            f"Photo size ({spec.width_mm}x{spec.height_mm}mm) does not fit "
            f"on {paper} paper at all. Choose a larger paper size."
        )

    grid_w = cols * photo_w_px + (cols - 1) * gutter_px
    grid_h = rows * photo_h_px + (rows - 1) * gutter_px
    origin_x = (sheet_w_px - grid_w) // 2
    origin_y = (sheet_h_px - grid_h) // 2

    sheet = Image.new("RGB", (sheet_w_px, sheet_h_px), "#FFFFFF")
    draw = ImageDraw.Draw(sheet)

    for r in range(rows):
        for c in range(cols):
            x = origin_x + c * (photo_w_px + gutter_px)
            y = origin_y + r * (photo_h_px + gutter_px)
            # Direct 1:1 paste — no resampling on the tile itself, so
            # zero distortion/stretch versus the already-correct source.
            sheet.paste(photo, (x, y))
            draw.rectangle(
                [x, y, x + photo_w_px - 1, y + photo_h_px - 1],
                outline=BORDER_HEX,
                width=BORDER_WIDTH_PX,
            )

    del draw
    gc.collect()
    return sheet


# ---------------------------------------------------------------------------
# Stage 6.5 — compliance heuristics
# ---------------------------------------------------------------------------
# Honest scope: these are geometric/pixel heuristics off YuNet's box, score,
# and 5-point landmarks — NOT a trained compliance classifier. They catch
# clear failures (face too small, extreme tilt, low detector confidence,
# bright glare in the eye region) but cannot verify things a real classifier
# would (eyes actually open vs. closed, genuine neutral expression, printed
# background uniformity beyond the color we composited ourselves). Every
# check below is labeled for what it actually measures — never claim a
# check we can't back, per the "no fake claims" rule this was built under.
@dataclass(frozen=True)
class ComplianceCheck:
    label: str
    passed: bool
    detail: str


def check_compliance(
    bounded: Image.Image, face: FaceBox, spec: "PhotoSpec"
) -> list[ComplianceCheck]:
    """Run cheap geometric/pixel heuristics on the pre-crop bounded image
    and detected face. Returns a list of pass/fail checks with plain-
    English detail for the UI to render as a checklist.
    """
    checks: list[ComplianceCheck] = []
    img_w, img_h = bounded.size

    # --- detector confidence ---
    # YuNet's own score is its confidence the box IS a face — low score
    # correlates with occlusion, extreme angle, or a false-positive match,
    # not with "photo quality" directly, but it's the most honest single
    # number the detector gives us.
    checks.append(
        ComplianceCheck(
            "Face detection confidence",
            face.score >= 0.85,
            f"{face.score * 100:.0f}% confidence"
            + ("" if face.score >= 0.85 else " — try a clearer, front-facing shot"),
        )
    )

    # --- face size relative to frame ---
    # Too-small a face in the source photo means the auto-crop has to
    # upscale heavily to hit the spec's head-ratio, softening detail.
    face_frac = (face.w * face.h) / (img_w * img_h)
    checks.append(
        ComplianceCheck(
            "Face size in source photo",
            face_frac >= 0.03,
            "Face fills enough of the frame"
            if face_frac >= 0.03
            else "Face is small in the source photo — move closer to the camera",
        )
    )

    # --- head tilt, via eye-line angle from landmarks ---
    # right_eye, left_eye are landmarks[0], landmarks[1]. A level eye-line
    # is the standard passport-photo "no head tilt" proxy every commercial
    # tool uses off a 2-point eye estimate — this is that same estimate,
    # not a full 3D pose model.
    tilt_ok = True
    tilt_detail = "Eye-line estimate unavailable"
    if face.landmarks and len(face.landmarks) >= 2:
        (rx, ry), (lx, ly) = face.landmarks[0], face.landmarks[1]
        import math

        angle_deg = abs(math.degrees(math.atan2(ly - ry, lx - rx)))
        tilt_ok = angle_deg <= 8.0
        tilt_detail = (
            f"Head level (~{angle_deg:.0f}° tilt)"
            if tilt_ok
            else f"Head appears tilted (~{angle_deg:.0f}°) — face the camera directly"
        )
    checks.append(ComplianceCheck("Head not tilted", tilt_ok, tilt_detail))

    # --- glare/glasses-glint proxy over the eye regions ---
    # Samples a small patch around each eye landmark and flags a large
    # cluster of near-pure-white pixels — a real proxy for lens glare, not
    # a "glasses detected" classifier. Photos with no glasses simply pass
    # this check trivially (no bright cluster to find).
    glare_ok = True
    glare_detail = "No strong glare detected near eyes"
    if face.landmarks and len(face.landmarks) >= 2:
        np_img = np.array(bounded.convert("RGB"))
        patch_r = max(4, int(face.w * 0.12))
        bright_frac_max = 0.0
        for (ex, ey) in face.landmarks[:2]:
            ex, ey = int(ex), int(ey)
            y0, y1 = max(0, ey - patch_r), min(img_h, ey + patch_r)
            x0, x1 = max(0, ex - patch_r), min(img_w, ex + patch_r)
            patch = np_img[y0:y1, x0:x1]
            if patch.size == 0:
                continue
            bright = np.all(patch > 235, axis=-1).mean()
            bright_frac_max = max(bright_frac_max, float(bright))
        del np_img
        glare_ok = bright_frac_max < 0.35
        if not glare_ok:
            glare_detail = "Possible glare on glasses/eyes — try removing glasses or adjusting lighting"
    checks.append(ComplianceCheck("No lens glare", glare_ok, glare_detail))

    return checks


def format_compliance_markdown(checks: list["ComplianceCheck"]) -> str:
    """Render checks as a compact markdown checklist for the Gradio UI."""
    lines = ["**Compliance check** (automated heuristics — always verify against official rules):"]
    for c in checks:
        icon = "✅" if c.passed else "⚠️"
        lines.append(f"- {icon} {c.label}: {c.detail}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Stage 5.5 — outfit overlay (garment compositing)
# ---------------------------------------------------------------------------
# How this works, honestly: this is composite-based garment overlay, the
# same technique cutout.pro's passport tool uses (confirmed by inspecting
# their garment assets — pre-rendered transparent PNGs cropped at the
# collar/shoulder line, not full-body generative reclothing). We detect
# the subject's shoulder keypoints, scale a pre-made garment PNG to match
# their shoulder width, and composite it over the torso region — under
# the face, over the original clothing. This is NOT clothing-aware
# (it won't preserve the person's actual shirt collar poking through a
# V-neck garment, for example) — it is a neck-down garment swap, which is
# exactly what passport-photo outfit tools need since only the shoulders-
# up region matters for the final crop.
try:
    import onnxruntime as ort
    _ONNXRUNTIME_AVAILABLE = True
except ImportError:
    _ONNXRUNTIME_AVAILABLE = False

_MOVENET_MODEL_ID = "Xenova/movenet-singlepose-lightning"
_MOVENET_FILENAME = "onnx/model.onnx"
_MOVENET_INPUT_SIZE = 192  # MoveNet Lightning's fixed input resolution —
# not configurable per the model architecture, unlike YuNet's setInputSize.

_pose_session_lock = threading.Lock()
_pose_session: Optional["ort.InferenceSession"] = None

# COCO-style 17 keypoint indices MoveNet outputs, in order. We only need
# shoulders, but documenting the full layout avoids future confusion if
# more keypoints (hips, for a future full-body feature) get used later.
_KP_LEFT_SHOULDER = 5
_KP_RIGHT_SHOULDER = 6


def _get_pose_session() -> "ort.InferenceSession":
    global _pose_session
    if not _ONNXRUNTIME_AVAILABLE:
        raise RuntimeError(
            "onnxruntime is not installed — outfit overlay requires it. "
            "Check requirements.txt."
        )
    with _pose_session_lock:
        if _pose_session is None:
            from huggingface_hub import hf_hub_download

            model_path = hf_hub_download(_MOVENET_MODEL_ID, _MOVENET_FILENAME)
            # CPUExecutionProvider deliberately — MoveNet Lightning is a
            # ~9MB model that runs in single-digit milliseconds on CPU;
            # routing it through the ZeroGPU @spaces.GPU machinery would
            # add GPU-attach overhead (seconds) for a task that doesn't
            # need it. Only BiRefNet's much heavier segmentation pass
            # (Stage 2) is worth the GPU round-trip.
            _pose_session = ort.InferenceSession(
                model_path, providers=["CPUExecutionProvider"]
            )
    return _pose_session


@dataclass(frozen=True)
class ShoulderKeypoints:
    left_x: float
    left_y: float
    right_x: float
    right_y: float
    confidence: float  # min of the two keypoint confidences

    @property
    def width_px(self) -> float:
        return abs(self.right_x - self.left_x)

    @property
    def center_x(self) -> float:
        return (self.left_x + self.right_x) / 2

    @property
    def center_y(self) -> float:
        return (self.left_y + self.right_y) / 2


def detect_shoulders(image_rgb: Image.Image) -> Optional[ShoulderKeypoints]:
    """Run MoveNet Lightning on the full bounded frame and return shoulder
    keypoints in the image's own pixel coordinates. Returns None (not a
    raised error) if confidence is too low — outfit overlay is an
    optional enhancement, so a low-confidence pose read should silently
    disable the feature for this photo rather than fail the whole
    pipeline the way a missing face does.
    """
    session = _get_pose_session()
    img = image_rgb.convert("RGB")
    orig_w, orig_h = img.size

    resized = img.resize(
        (_MOVENET_INPUT_SIZE, _MOVENET_INPUT_SIZE), Image.Resampling.BILINEAR
    )
    # MoveNet's published input contract: int32 tensor, NHWC, [0,255] raw
    # pixel values (no normalization) — this is the model's own expected
    # format, not a convention we chose.
    inp = np.array(resized, dtype=np.int32)[np.newaxis, ...]
    del resized

    outputs = session.run(None, {session.get_inputs()[0].name: inp})
    del inp
    # Output shape (1,1,17,3): [y, x, confidence] per keypoint, normalized
    # to [0,1] against the model's own 192x192 input frame.
    keypoints = outputs[0][0, 0]
    del outputs
    gc.collect()

    ly, lx, lc = keypoints[_KP_LEFT_SHOULDER]
    ry, rx, rc = keypoints[_KP_RIGHT_SHOULDER]
    confidence = float(min(lc, rc))

    if confidence < 0.3:
        return None

    return ShoulderKeypoints(
        left_x=float(lx) * orig_w,
        left_y=float(ly) * orig_h,
        right_x=float(rx) * orig_w,
        right_y=float(ry) * orig_h,
        confidence=confidence,
    )


# ---------------------------------------------------------------------------
# Garment asset registry
# ---------------------------------------------------------------------------
_GARMENTS_DIR = os.path.join(os.path.dirname(__file__), "assets", "garments")


@dataclass(frozen=True)
class GarmentAsset:
    garment_id: str
    label: str
    image: Image.Image  # pre-loaded RGBA, cached for process lifetime
    shoulder_width_px: int
    shoulder_y_px: int
    shoulder_cx_px: int
    collar_y_px: int
    collar_cx_px: int


_garment_cache: dict[str, GarmentAsset] = {}
_garment_cache_lock = threading.Lock()

# Display label per garment_id — kept separate from the filename/id so the
# UI can show something human-friendly without renaming asset files.
GARMENT_LABELS: dict[str, str] = {
    "mens_navy_suit_tie": "Men's Navy Suit + Tie",
    "mens_navy_suit_pocket_square": "Men's Navy Suit + Pocket Square",
    "womens_blush_blazer": "Women's Blush Blazer",
}


def list_garments() -> list[str]:
    """Return available garment display labels, discovered from whatever
    normalized .png/.json pairs exist in assets/garments/ — so dropping in
    a new pair (via normalize_garments.py) makes it available without a
    code change here.
    """
    if not os.path.isdir(_GARMENTS_DIR):
        return []
    ids = sorted(
        f[:-5] for f in os.listdir(_GARMENTS_DIR) if f.endswith(".json")
    )
    return [GARMENT_LABELS.get(gid, gid) for gid in ids]


def _label_to_id(label: str) -> Optional[str]:
    for gid, lbl in GARMENT_LABELS.items():
        if lbl == label:
            return gid
    # Fallback: label IS the id (covers any garment dropped in without a
    # GARMENT_LABELS entry — list_garments() would have returned the raw
    # id as its own label in that case).
    if os.path.exists(os.path.join(_GARMENTS_DIR, f"{label}.json")):
        return label
    return None


def _load_garment(garment_id: str) -> GarmentAsset:
    with _garment_cache_lock:
        if garment_id in _garment_cache:
            return _garment_cache[garment_id]

        json_path = os.path.join(_GARMENTS_DIR, f"{garment_id}.json")
        png_path = os.path.join(_GARMENTS_DIR, f"{garment_id}.png")
        if not (os.path.exists(json_path) and os.path.exists(png_path)):
            raise ValueError(f"Unknown garment: {garment_id}")

        import json

        with open(json_path) as f:
            anchor = json.load(f)

        img = Image.open(png_path).convert("RGBA")
        asset = GarmentAsset(
            garment_id=garment_id,
            label=GARMENT_LABELS.get(garment_id, garment_id),
            image=img,
            shoulder_width_px=anchor["shoulder_width_px"],
            # Fallback to collar_y/collar_cx for any garment normalized
            # before shoulder_y_px/shoulder_cx_px were split out as
            # separate fields — keeps old JSON sidecars from hard-erroring,
            # though re-running normalize_garments.py is the real fix.
            shoulder_y_px=anchor.get("shoulder_y_px", anchor["collar_y_px"]),
            shoulder_cx_px=anchor.get("shoulder_cx_px", anchor["collar_cx_px"]),
            collar_y_px=anchor["collar_y_px"],
            collar_cx_px=anchor["collar_cx_px"],
        )
        _garment_cache[garment_id] = asset
        return asset


# Fallback ratio used when live shoulder detection fails/low-confidence:
# garment shoulder-width as a multiple of face width. Derived from typical
# adult head-to-shoulder proportions (shoulder span ~= 2.2-2.6x face
# width for a frontal passport-style pose) — a reasonable default, not a
# substitute for the real pose read when it's available.
_FALLBACK_SHOULDER_TO_FACE_RATIO = 2.4


def apply_outfit(
    bounded: Image.Image,
    matted: Image.Image,
    face: "FaceBox",
    spec: "PhotoSpec",
    garment_label: str,
    zoom: float = 1.0,
    x_offset: float = 0.0,
    y_offset: float = 0.0,
) -> Image.Image:
    """Composite the chosen garment onto `matted` (the alpha-matted
    subject), scaled to fit the space that will actually survive into the
    final cropped photo.

    Why this needs `spec` (and the same zoom/offset params as
    crop_to_spec): a passport photo's crop window shows only a small
    sliver below the chin — typically 40–100px in source-image terms,
    far less than a garment scaled to match real shoulder width would
    need. Earlier versions of this function scaled the garment to match
    MoveNet's detected shoulder width, which produced a garment 2–3x
    taller than the space available below the chin — confirmed via
    runtime logs showing e.g. a 141px-tall garment against a ~50px
    available strip, so almost all of it was cropped away regardless of
    vertical anchor position. The fix: compute the SAME crop box
    crop_to_spec() will use, measure how much vertical space exists
    between the chin and the crop's bottom edge, and scale the garment's
    HEIGHT to fill that space (preserving aspect ratio) — not its width
    to match shoulder measurements. This guarantees the garment's visible
    portion actually reaches the crop boundary instead of being a tiny
    cropped-off fragment.

    Returns a NEW RGBA image — does not mutate `matted` in place, so the
    caller's reference to the pre-outfit matte stays valid if needed
    elsewhere (e.g. the bg_removed_preview stage thumbnail should show
    the ORIGINAL matte, not the outfit-composited one).

    Garment is composited BELOW the face region — we paste the garment
    layer first, then paste the ORIGINAL matted subject's head/face
    region back on top, so the person's real face is never occluded by
    the garment PNG.
    """
    garment_id = _label_to_id(garment_label)
    if garment_id is None:
        raise ValueError(f"Unknown garment: {garment_label}")
    garment = _load_garment(garment_id)

    img_w, img_h = matted.size
    crop_left, crop_top, crop_right, crop_bottom = compute_crop_box(
        img_w, img_h, face, spec, zoom, x_offset, y_offset
    )

    chin_y = face.y + face.h * 1.35
    # Available vertical space between the chin and the bottom of the
    # crop window — this, not shoulder width, is what the garment must
    # be scaled to fill. Guard against a degenerate/negative value (an
    # extreme manual zoom/offset could in principle push crop_bottom
    # above chin_y) with a small sane floor.
    available_height = max(8.0, crop_bottom - chin_y)

    # Horizontal position still uses shoulder detection when confident —
    # this determines WHERE (left-right) the garment centers, not how
    # large it is. Falls back to face-width-derived center when pose
    # confidence is low.
    shoulders = detect_shoulders(bounded)
    if shoulders is not None and shoulders.confidence >= 0.3:
        target_cx = shoulders.center_x
        _outfit_debug_source = "pose"
    else:
        target_cx = face.cx
        _outfit_debug_source = "fallback"

    # Scale by HEIGHT to fill the available strip below the chin, with a
    # small overshoot factor so the garment's edge runs slightly past the
    # crop boundary rather than leaving a visible gap if this estimate is
    # a little short — the face-reinstate patch above the garment and
    # the crop boundary below it both hide any resulting overshoot.
    overshoot = 1.15
    scale = (available_height * overshoot) / garment.image.height
    new_w = max(1, int(round(garment.image.width * scale)))
    new_h = max(1, int(round(garment.image.height * scale)))
    scaled_garment = garment.image.resize((new_w, new_h), Image.Resampling.LANCZOS)

    # Horizontal: center the scaled garment's own collar point at
    # target_cx. Vertical: place the garment's collar point at the chin
    # (small gap below it) — same anchor concept as before, but now the
    # garment's overall size is correct for the space it needs to fill.
    scaled_collar_x = garment.collar_cx_px * scale
    scaled_collar_y = garment.collar_y_px * scale
    collar_gap_px = face.h * 0.15
    target_collar_y = chin_y + collar_gap_px
    paste_x = int(round(target_cx - scaled_collar_x))
    paste_y = int(round(target_collar_y - scaled_collar_y))

    logging.getLogger("passport-maker").info(
        "apply_outfit: source=%s shoulders_conf=%.2f chin_y=%.0f "
        "crop_box=(%.0f,%.0f,%.0f,%.0f) available_height=%.0f scale=%.3f "
        "garment_size=%dx%d target_cx=%.0f target_collar_y=%.0f "
        "paste=(%d,%d) canvas=%dx%d",
        _outfit_debug_source,
        shoulders.confidence if shoulders else -1.0,
        chin_y, crop_left, crop_top, crop_right, crop_bottom,
        available_height, scale, new_w, new_h, target_cx, target_collar_y,
        paste_x, paste_y, matted.width, matted.height,
    )

    # Composite: start from a copy of matted, paste garment on top (using
    # its own alpha as the mask so transparent garment-PNG pixels don't
    # overwrite the subject), THEN paste the original head/shoulders
    # region from `matted` back on top of that — guarantees the face is
    # never covered by garment pixels regardless of alignment error.
    result = matted.copy()
    result.paste(scaled_garment, (paste_x, paste_y), scaled_garment)
    del scaled_garment
    gc.collect()

    # Re-apply the original face region on top — but ONLY the face
    # itself, not a large margin around it. An elliptical mask, sized to
    # face size with only a small margin, keeps this from reaching down
    # into the (now much closer, since the garment is properly sized)
    # collar area.
    pad = int(max(face.w, face.h) * 0.15)
    side = max(face.w, face.h) + 2 * pad
    fx0 = max(0, int(face.cx - side / 2))
    fy0 = max(0, int(face.cy - side / 2))
    fx1 = min(matted.width, fx0 + side)
    fy1 = min(matted.height, fy0 + side)
    face_patch = matted.crop((fx0, fy0, fx1, fy1))

    from PIL import ImageDraw as _ImageDraw
    import numpy as _np

    ellipse_mask = Image.new("L", face_patch.size, 0)
    _ImageDraw.Draw(ellipse_mask).ellipse([0, 0, face_patch.size[0], face_patch.size[1]], fill=255)

    if face_patch.mode == "RGBA":
        orig_alpha = face_patch.split()[3]
        combined_arr = (
            _np.array(ellipse_mask, dtype=_np.uint16)
            * _np.array(orig_alpha, dtype=_np.uint16)
            // 255
        ).astype(_np.uint8)
        combined_mask = Image.fromarray(combined_arr, mode="L")
    else:
        combined_mask = ellipse_mask

    result.paste(face_patch, (fx0, fy0), combined_mask)
    del face_patch, ellipse_mask, combined_mask
    gc.collect()

    return result


# ---------------------------------------------------------------------------
# Orchestration — the single entry point ui.py calls
# ---------------------------------------------------------------------------
def process_photo(
    image: Image.Image,
    spec_key: str,
    bg_hex: Optional[str],
    paper_key: Optional[str],
    zoom: float = 1.0,
    x_offset: float = 0.0,
    y_offset: float = 0.0,
    auto_straighten: bool = True,
    outfit_label: Optional[str] = None,
) -> tuple[Image.Image, Optional[Image.Image], list["ComplianceCheck"], Image.Image, Image.Image, float, bool, Optional[str]]:
    """Full pipeline. Returns (single_photo, print_sheet_or_None,
    compliance_checks, bg_removed_preview, face_only_thumbnail,
    straighten_angle_applied, outfit_applied, outfit_error).

    bg_removed_preview: the alpha-matted subject on transparent background,
    at the same size as `bounded` — this is a display artifact for the UI's
    stage-by-stage view (mirrors what cutout.pro shows as its "Result"
    step), not used further in the pipeline itself. Shows the ORIGINAL
    matte even when outfit_label is set — outfit compositing is a
    downstream step, not part of what "background removed" should depict.

    face_only_thumbnail: a tight square crop around the detected face,
    also transparent-background — display-only, same purpose.

    straighten_angle_applied: degrees the whole frame was rotated to level
    the eye-line (0.0 if auto_straighten=False or tilt was negligible).
    Head AND torso/shoulders straighten together because whole-frame
    rotation moves every pixel in the rigid photo by the same amount —
    see straighten_image()'s docstring for why this is the physically
    correct operation for a single-subject photo, not a head-only crop.

    outfit_label: display label of a garment from list_garments(), or
    None/"" to skip outfit overlay entirely (default — the original
    photo's clothing is used, exactly as before this feature existed).

    outfit_applied: True only if outfit compositing genuinely succeeded.
    False whenever outfit_label was set but compositing failed and the
    pipeline silently fell back to the original photo — the caller MUST
    check this rather than assuming outfit_label being set means the
    photo was actually outfitted, since that assumption previously
    produced a status message claiming an outfit was applied when it
    silently wasn't.

    outfit_error: short error string when outfit_applied is False due to
    a failure (None if no outfit was requested, or if it succeeded).

    Raises ValueError with a user-facing message on any recoverable
    failure (no face found, bad spec key, etc) — ui.py surfaces these via
    gr.Error rather than letting a raw traceback reach the user.
    """
    if spec_key not in STANDARDS:
        raise ValueError(f"Unknown photo standard: {spec_key}")
    spec = STANDARDS[spec_key]
    effective_bg = bg_hex or spec.bg_hex

    bounded = load_and_bound(image)

    # Face detection MUST run on the same pixel grid as the alpha matte
    # (both derive from `bounded`), otherwise the face box coordinates
    # used in crop_to_spec() would be wrong-scale against the matted
    # image and produce a badly-centered crop.
    face = detect_main_face(bounded)

    straighten_angle_applied = 0.0
    if auto_straighten:
        tilt_angle = compute_tilt_angle(face)
        if abs(tilt_angle) >= 0.5:
            straightened = straighten_image(bounded, tilt_angle)
            if straightened is not bounded:  # rotation actually happened
                del bounded
                bounded = straightened
                # Coordinates changed (rotation + expand=True resized the
                # canvas) — re-detect rather than analytically transform
                # the old box, matching straighten_image()'s docstring.
                face = detect_main_face(bounded)
                straighten_angle_applied = max(
                    -MAX_AUTO_STRAIGHTEN_DEG, min(MAX_AUTO_STRAIGHTEN_DEG, tilt_angle)
                )

    # Compliance heuristics need the pre-matte `bounded` pixels (glare
    # check reads real image brightness) and `face` — must run before
    # `bounded` is freed below. Runs on the already-straightened frame so
    # the tilt check reflects the final, corrected state.
    checks = check_compliance(bounded, face, spec)

    matted = segment_alpha(bounded)

    # Outfit overlay needs `bounded` (real RGB pixels, for MoveNet's pose
    # read) — must run before `bounded` is freed below. If no outfit was
    # requested, skip entirely: zero cost, zero behavior change from
    # before this feature existed.
    outfitted = None
    outfit_applied = False
    outfit_error: Optional[str] = None
    if outfit_label:
        try:
            outfitted = apply_outfit(bounded, matted, face, spec, outfit_label, zoom=zoom, x_offset=x_offset, y_offset=y_offset)
            outfit_applied = True
        except ValueError:
            raise  # unknown garment label — genuine user-facing error
        except Exception as e:
            # Pose detection or compositing failed for a reason that
            # isn't the user's fault (e.g. onnxruntime hiccup) — degrade
            # gracefully to the original photo rather than failing the
            # whole generate. Outfit overlay is an enhancement, not a
            # core guarantee the way face detection is. BUT: log it for
            # real, and tell the caller it silently degraded — an earlier
            # version of this code swallowed the exception AND still
            # reported "outfit: X" in the UI status line, which lied to
            # the user about what actually happened to their photo.
            import traceback
            logging.getLogger("passport-maker").warning(
                "Outfit overlay failed, falling back to original photo: %s",
                traceback.format_exc(),
            )
            outfitted = None
            outfit_applied = False
            outfit_error = str(e) or type(e).__name__

    del bounded
    gc.collect()

    # --- display-only face thumbnail, built from the ORIGINAL `matted`
    # before it's consumed by crop_to_spec() below. Square crop, generous
    # padding around the detected box so the thumbnail reads as "a
    # headshot," not a tight bounding-box rectangle. Clamped to image
    # bounds — no padding added here (unlike crop_to_spec) since this is
    # a preview, not a spec-exact deliverable.
    pad = int(max(face.w, face.h) * 0.6)
    side = max(face.w, face.h) + 2 * pad
    fx0 = max(0, int(face.cx - side / 2))
    fy0 = max(0, int(face.cy - side / 2))
    fx1 = min(matted.width, fx0 + side)
    fy1 = min(matted.height, fy0 + side)
    face_thumb = matted.crop((fx0, fy0, fx1, fy1))

    bg_removed_preview = matted.copy()

    # Crop the outfitted version if one was produced, otherwise fall back
    # to the original matte — this is the only place the two diverge, so
    # everything downstream (crop/composite/print-sheet) is identical
    # code regardless of whether an outfit was applied.
    source_for_crop = outfitted if outfitted is not None else matted
    cropped = crop_to_spec(source_for_crop, face, spec, zoom=zoom, x_offset=x_offset, y_offset=y_offset)
    del matted
    if outfitted is not None:
        del outfitted
    gc.collect()

    final_photo = composite_background(cropped, effective_bg)
    del cropped
    gc.collect()

    sheet = None
    if paper_key:
        sheet = build_print_sheet(final_photo, spec, paper_key)

    return final_photo, sheet, checks, bg_removed_preview, face_thumb, straighten_angle_applied, outfit_applied, outfit_error


MAX_BATCH_SIZE = 10  # cap: each image is a separate @spaces.GPU acquisition
# (up to 20s A10G time requested per call) — an unbounded batch from one
# submit could starve the shared ZeroGPU queue for other users. 10 images
# at worst-case ~6-8s/image on CPU-fallback is also a sane wall-clock
# ceiling before a browser tab feels "stuck".


@dataclass(frozen=True)
class BatchResult:
    filename: str
    photo: Optional[Image.Image]
    error: Optional[str]


def process_batch(
    images: list[tuple[str, Image.Image]],
    spec_key: str,
    bg_hex: Optional[str],
    zoom: float = 1.0,
    x_offset: float = 0.0,
    y_offset: float = 0.0,
    auto_straighten: bool = True,
    outfit_label: Optional[str] = None,
) -> list[BatchResult]:
    """Run process_photo across multiple images. Never lets one bad image
    (no face detected, corrupt file, etc) abort the whole batch — each
    failure is captured per-item so the user gets N-1 good results instead
    of a single error wiping everything. No print-sheet tiling in batch
    mode (each photo is a different source person/crop; tiling assumes one
    subject repeated, which does not apply here).
    """
    if len(images) > MAX_BATCH_SIZE:
        raise ValueError(
            f"Batch limit is {MAX_BATCH_SIZE} photos per submission. "
            f"You uploaded {len(images)} — please split into smaller batches."
        )

    results: list[BatchResult] = []
    for filename, img in images:
        try:
            photo, _sheet, _checks, _bg_preview, _face_thumb, _angle, _outfit_ok, _outfit_err = process_photo(
                image=img,
                spec_key=spec_key,
                bg_hex=bg_hex,
                paper_key=None,
                zoom=zoom,
                x_offset=x_offset,
                y_offset=y_offset,
                auto_straighten=auto_straighten,
                outfit_label=outfit_label,
            )
            results.append(BatchResult(filename, photo, None))
        except ValueError as e:
            results.append(BatchResult(filename, None, str(e)))
        except Exception:
            results.append(
                BatchResult(filename, None, "Processing failed — try a different photo.")
            )
        gc.collect()
    return results