Update engine.py

engine.py

@@ -159,6 +159,21 @@ def warm_up() -> None:
     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
@@ -916,6 +931,307 @@ def format_compliance_markdown(checks: list["ComplianceCheck"]) -> str:
     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
+    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"],
+            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",
+    garment_label: str,
+) -> Image.Image:
+    """Composite the chosen garment onto `matted` (the alpha-matted
+    subject), aligned to detected shoulder keypoints when confidently
+    available, falling back to a face-width-derived estimate otherwise.
+
+    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, matching what
+    cutout.pro's own stage breakdown shows).
+
+    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 even if the garment's collar extends higher than the
+    detected shoulder line.
+    """
+    garment_id = _label_to_id(garment_label)
+    if garment_id is None:
+        raise ValueError(f"Unknown garment: {garment_label}")
+    garment = _load_garment(garment_id)
+
+    shoulders = detect_shoulders(bounded)
+
+    if shoulders is not None and shoulders.confidence >= 0.3:
+        target_width_px = shoulders.width_px
+        target_cx = shoulders.center_x
+        # Collar should sit at the shoulder line, not above/below it —
+        # use the shoulder keypoints' own y as the collar target.
+        target_collar_y = shoulders.center_y
+    else:
+        # Fallback: derive an approximate shoulder position from the
+        # already-known face box, since we always have that.
+        target_width_px = face.w * _FALLBACK_SHOULDER_TO_FACE_RATIO
+        target_cx = face.cx
+        # Shoulders sit below the chin by roughly one more face-height —
+        # same order-of-magnitude approximation used elsewhere in this
+        # file for anatomy without a direct measurement.
+        target_collar_y = face.y + face.h * 1.9
+
+    scale = target_width_px / garment.shoulder_width_px
+    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)
+
+    # Position so the garment's own collar anchor point lands exactly at
+    # (target_cx, target_collar_y) in the destination image.
+    scaled_collar_x = garment.collar_cx_px * scale
+    scaled_collar_y = garment.collar_y_px * scale
+    paste_x = int(round(target_cx - scaled_collar_x))
+    paste_y = int(round(target_collar_y - scaled_collar_y))
+
+    # 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. Padding matches the
+    # face-thumbnail crop elsewhere in this file (0.6x face size margin)
+    # so the reinstated patch comfortably covers the whole head with a
+    # soft-enough boundary that a hard rectangle edge is unlikely to fall
+    # across a hairline in a way that reads as an obvious seam.
+    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_patch = matted.crop((fx0, fy0, fx1, fy1))
+    result.paste(face_patch, (fx0, fy0), face_patch)
+    del face_patch
+    gc.collect()
+
+    return result
+
+
 # ---------------------------------------------------------------------------
 # Orchestration — the single entry point ui.py calls
 # ---------------------------------------------------------------------------
@@ -928,6 +1244,7 @@ def process_photo(
     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]:
     """Full pipeline. Returns (single_photo, print_sheet_or_None,
     compliance_checks, bg_removed_preview, face_only_thumbnail,
@@ -936,7 +1253,9 @@ def process_photo(
     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.
+    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.
@@ -948,6 +1267,10 @@ def process_photo(
     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).
+
     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.
@@ -988,15 +1311,34 @@ def process_photo(
     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
+    if outfit_label:
+        try:
+            outfitted = apply_outfit(bounded, matted, face, outfit_label)
+        except ValueError:
+            raise  # unknown garment label — genuine user-facing error
+        except Exception:
+            # 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.
+            outfitted = None
+
     del bounded
     gc.collect()
 
-    # --- display-only face thumbnail, built from `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.
+    # --- 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))
@@ -1007,8 +1349,15 @@ def process_photo(
 
     bg_removed_preview = matted.copy()
 
-    cropped = crop_to_spec(matted, face, spec, zoom=zoom, x_offset=x_offset, y_offset=y_offset)
+    # 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)
@@ -1044,6 +1393,7 @@ def process_batch(
     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
@@ -1070,6 +1420,7 @@ def process_batch(
                 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: