app.py

"""
ui.py — Gradio Blocks frontend for the Passport Photo Maker.

Entry point for the HF Space (Gradio SDK auto-runs this file as `app.py`).
"""

from __future__ import annotations

import gc
import io
import logging
import os
import tempfile
import zipfile

import gradio as gr
from PIL import Image

from engine import (
    STANDARDS,
    PAPER_SIZES_MM,
    MAX_BATCH_SIZE,
    process_photo,
    process_batch,
    format_compliance_markdown,
    list_garments,
    warm_up,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("passport-maker")

STANDARD_CHOICES = list(STANDARDS.keys())
PAPER_CHOICES = ["None (single photo only)"] + list(PAPER_SIZES_MM.keys())
OUTFIT_CHOICES = ["None (keep original clothing)"] + list_garments()


# ---------------------------------------------------------------------------
# Single-photo handler
# ---------------------------------------------------------------------------
def run_pipeline(
    image: Image.Image | None,
    spec_choice: str,
    bg_color: str | None,
    paper_choice: str,
    zoom: float,
    x_offset: float,
    y_offset: float,
    auto_straighten: bool,
    outfit_choice: str,
    progress: gr.Progress = gr.Progress(),
):
    """Gradio click handler. Wraps engine.process_photo with progress
    updates and converts internal exceptions into user-facing gr.Error
    popups instead of raw tracebacks.
    """
    if image is None:
        raise gr.Error("Please upload a photo first.")

    try:
        progress(0.15, desc="Preparing image...")
        progress(0.35, desc="Removing background (this takes a few seconds on CPU)...")

        paper_key = None if paper_choice.startswith("None") else paper_choice
        outfit_label = None if outfit_choice.startswith("None") else outfit_choice

        photo, sheet, checks, bg_removed_preview, face_thumb, straighten_angle, outfit_applied, outfit_error = process_photo(
            image=image,
            spec_key=spec_choice,
            bg_hex=bg_color,
            paper_key=paper_key,
            zoom=zoom,
            x_offset=x_offset,
            y_offset=y_offset,
            auto_straighten=auto_straighten,
            outfit_label=outfit_label,
        )

        progress(0.9, desc="Finalizing...")

        gc.collect()
        progress(1.0, desc="Done")

        status = f"✅ **{spec_choice}** — {STANDARDS[spec_choice].width_mm}×{STANDARDS[spec_choice].height_mm}mm @ 300 DPI ({photo.width}×{photo.height}px)"
        if sheet is not None:
            status += f" · tiled on {paper_key}"
        if abs(straighten_angle) >= 0.5:
            status += f" · auto-straightened {abs(straighten_angle):.1f}°"
        if outfit_label:
            # Only claim the outfit was applied if it genuinely was — a
            # prior version of this code claimed success here whenever
            # outfit_label was merely selected, even when compositing had
            # silently failed and fallen back to the original photo. Now
            # honest either way.
            if outfit_applied:
                status += f" · outfit: {outfit_label}"
            else:
                status += (
                    f" · ⚠️ outfit '{outfit_label}' could not be applied "
                    f"(showing original clothing)"
                )
                if outfit_error:
                    status += f" — {outfit_error}"

        compliance_md = format_compliance_markdown(checks)

        # Stage-by-stage gallery — mirrors cutout.pro's Original / BG-removed
        # / Face cutout / Result breakdown. Each tuple is (image, caption);
        # Gradio Gallery renders the caption under the thumbnail.
        stages = [
            (image, "1. Original"),
            (bg_removed_preview, "2. Background Removed"),
            (face_thumb, "3. Face Cutout"),
            (photo, "4. Cropped / Final"),
        ]

        # ImageSlider needs a (before, after) pair at the same canvas. We
        # hand it (original upload, final result) — Gradio letterboxes/
        # fits internally, so mismatched source aspect ratios display fine;
        # it's a visual compare, not a pixel-aligned overlay.
        return stages, (image, photo), sheet, status, compliance_md

    except ValueError as e:
        # Expected, user-facing errors (no face detected, bad spec, photo
        # doesn't fit paper, etc) — clean message, no traceback.
        raise gr.Error(str(e))
    except Exception:
        logger.exception("Unhandled error in process_photo")
        raise gr.Error(
            "Something went wrong processing this photo. Try a different "
            "image, or a smaller file size."
        )


# ---------------------------------------------------------------------------
# Batch handler
# ---------------------------------------------------------------------------
def run_batch(
    files: list | None,
    spec_choice: str,
    bg_color: str | None,
    zoom: float,
    x_offset: float,
    y_offset: float,
    auto_straighten: bool,
    outfit_choice: str,
    progress: gr.Progress = gr.Progress(),
):
    """Gradio click handler for the Batch tab. Accepts a list of uploaded
    file paths (from gr.File multiple), runs each through process_photo,
    and returns a gallery of successes + a status report + a zip download.
    Per-image failures never abort the batch — see engine.process_batch.
    """
    if not files:
        raise gr.Error("Please upload at least one photo.")

    if len(files) > MAX_BATCH_SIZE:
        raise gr.Error(
            f"Batch limit is {MAX_BATCH_SIZE} photos. You uploaded {len(files)} "
            f"— please remove some and try again."
        )

    images: list[tuple[str, Image.Image]] = []
    for f in files:
        path = f.name if hasattr(f, "name") else f
        try:
            img = Image.open(path)
            img.load()
            images.append((os.path.basename(path), img))
        except Exception:
            images.append((os.path.basename(path), None))

    progress(0.1, desc=f"Processing {len(images)} photos...")

    # Filter out unreadable files up front with a clear per-file error,
    # rather than letting them crash into process_batch's image pipeline.
    valid = [(name, img) for name, img in images if img is not None]
    bad_names = [name for name, img in images if img is None]

    outfit_label = None if outfit_choice.startswith("None") else outfit_choice

    try:
        results = process_batch(
            valid,
            spec_key=spec_choice,
            bg_hex=bg_color,
            zoom=zoom,
            x_offset=x_offset,
            y_offset=y_offset,
            auto_straighten=auto_straighten,
            outfit_label=outfit_label,
        )
    except ValueError as e:
        raise gr.Error(str(e))
    except Exception:
        logger.exception("Unhandled error in process_batch")
        raise gr.Error("Something went wrong processing this batch.")

    progress(0.85, desc="Packaging results...")

    gallery_items = []
    ok_count = 0
    status_lines = []
    zip_path = None

    good_results = [r for r in results if r.photo is not None]
    if good_results:
        tmp_dir = tempfile.mkdtemp(prefix="passport_batch_")
        zip_path = os.path.join(tmp_dir, "passport_photos.zip")
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for r in good_results:
                buf = io.BytesIO()
                r.photo.save(buf, format="PNG")
                out_name = os.path.splitext(r.filename)[0] + "_passport.png"
                zf.writestr(out_name, buf.getvalue())
                gallery_items.append((r.photo, r.filename))
                ok_count += 1

    for r in results:
        if r.photo is not None:
            status_lines.append(f"✅ {r.filename}")
        else:
            status_lines.append(f"❌ {r.filename} — {r.error}")
    for name in bad_names:
        status_lines.append(f"❌ {name} — Could not read this file as an image.")

    header = f"**{ok_count}/{len(images)} photos processed successfully.**"
    status_md = header + "\n\n" + "\n".join(status_lines)

    gc.collect()
    progress(1.0, desc="Done")

    return gallery_items, status_md, zip_path


def build_interface() -> gr.Blocks:
    with gr.Blocks(
        title="Passport Photo Maker — Free, Instant, 300 DPI",
    ) as demo:
        gr.Markdown(
            """
            # 🛂 Passport Photo Maker
            AI background removal + auto face-centering + exact 300 DPI passport,
            visa, and ID photo sizing. Free. No login. No watermark.
            """
        )

        with gr.Tabs():
            # -----------------------------------------------------------
            # TAB 1 — Single photo
            # -----------------------------------------------------------
            with gr.Tab("Single Photo"):
                with gr.Row():
                    with gr.Column(scale=1):
                        inp_image = gr.Image(
                            type="pil",
                            label="Upload your photo",
                            sources=["upload", "webcam"],
                        )
                        spec_dropdown = gr.Dropdown(
                            choices=STANDARD_CHOICES,
                            value=STANDARD_CHOICES[0],
                            label="Photo Standard",
                        )
                        bg_color = gr.ColorPicker(
                            value="#1E3A8A",
                            label="Background Color (overrides standard default)",
                        )
                        paper_dropdown = gr.Dropdown(
                            choices=PAPER_CHOICES,
                            value=PAPER_CHOICES[0],
                            label="Print Sheet Layout",
                        )
                        outfit_dropdown = gr.Dropdown(
                            choices=OUTFIT_CHOICES,
                            value=OUTFIT_CHOICES[0],
                            label="Outfit (replaces clothing below the neck)",
                        )
                        with gr.Accordion("Adjust Crop (optional — auto-crop is usually correct)", open=False):
                            zoom_slider = gr.Slider(
                                minimum=0.5, maximum=2.0, value=1.0, step=0.05,
                                label="Zoom (higher = tighter crop on face)",
                            )
                            x_slider = gr.Slider(
                                minimum=-0.5, maximum=0.5, value=0.0, step=0.02,
                                label="Shift Left / Right",
                            )
                            y_slider = gr.Slider(
                                minimum=-0.5, maximum=0.5, value=0.0, step=0.02,
                                label="Shift Up / Down",
                            )
                            reset_crop_btn = gr.Button("Reset Crop", size="sm")
                        straighten_checkbox = gr.Checkbox(
                            value=True,
                            label="Auto-Straighten (levels head + shoulders if the photo is tilted)",
                        )
                        generate_btn = gr.Button("Generate Photo", variant="primary", size="lg")

                    with gr.Column(scale=1):
                        out_stages = gr.Gallery(
                            label="Processing Stages",
                            columns=4,
                            object_fit="contain",
                            height=200,
                        )
                        out_compare = gr.ImageSlider(
                            label="Before / After — drag to compare",
                            type="pil",
                        )
                        status_box = gr.Markdown()
                        compliance_box = gr.Markdown()
                        out_sheet = gr.Image(type="pil", label="Print Sheet")

                # Update the color picker default whenever the standard
                # changes, so the user sees the correct expected background
                # before they touch anything — manual override still wins.
                def _sync_bg_default(spec_choice: str):
                    return gr.update(value=STANDARDS[spec_choice].bg_hex)

                spec_dropdown.change(
                    fn=_sync_bg_default, inputs=spec_dropdown, outputs=bg_color
                )

                generate_btn.click(
                    fn=run_pipeline,
                    inputs=[
                        inp_image, spec_dropdown, bg_color, paper_dropdown,
                        zoom_slider, x_slider, y_slider, straighten_checkbox,
                        outfit_dropdown,
                    ],
                    outputs=[out_stages, out_compare, out_sheet, status_box, compliance_box],
                )

                reset_crop_btn.click(
                    fn=lambda: (gr.update(value=1.0), gr.update(value=0.0), gr.update(value=0.0)),
                    inputs=None,
                    outputs=[zoom_slider, x_slider, y_slider],
                )

            # -----------------------------------------------------------
            # TAB 2 — Batch
            # -----------------------------------------------------------
            with gr.Tab("Batch (up to 10 photos)"):
                gr.Markdown(
                    "Upload multiple photos, apply one standard + background to all, "
                    "download every result in a single ZIP. Each photo is processed "
                    "independently — one bad photo won't stop the rest."
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        batch_files = gr.File(
                            label="Upload photos",
                            file_count="multiple",
                            file_types=["image"],
                        )
                        batch_spec_dropdown = gr.Dropdown(
                            choices=STANDARD_CHOICES,
                            value=STANDARD_CHOICES[0],
                            label="Photo Standard (applied to all)",
                        )
                        batch_bg_color = gr.ColorPicker(
                            value="#1E3A8A",
                            label="Background Color (overrides standard default)",
                        )
                        with gr.Accordion("Adjust Crop (applied to all photos)", open=False):
                            batch_zoom_slider = gr.Slider(
                                minimum=0.5, maximum=2.0, value=1.0, step=0.05,
                                label="Zoom",
                            )
                            batch_x_slider = gr.Slider(
                                minimum=-0.5, maximum=0.5, value=0.0, step=0.02,
                                label="Shift Left / Right",
                            )
                            batch_y_slider = gr.Slider(
                                minimum=-0.5, maximum=0.5, value=0.0, step=0.02,
                                label="Shift Up / Down",
                            )

                        def _sync_batch_bg_default(spec_choice: str):
                            return gr.update(value=STANDARDS[spec_choice].bg_hex)

                        batch_spec_dropdown.change(
                            fn=_sync_batch_bg_default,
                            inputs=batch_spec_dropdown,
                            outputs=batch_bg_color,
                        )

                        batch_straighten_checkbox = gr.Checkbox(
                            value=True,
                            label="Auto-Straighten (applied to all photos)",
                        )
                        batch_outfit_dropdown = gr.Dropdown(
                            choices=OUTFIT_CHOICES,
                            value=OUTFIT_CHOICES[0],
                            label="Outfit (applied to all photos)",
                        )
                        batch_generate_btn = gr.Button(
                            "Generate All", variant="primary", size="lg"
                        )

                    with gr.Column(scale=1):
                        batch_gallery = gr.Gallery(
                            label="Results", columns=3, object_fit="contain", height="auto"
                        )
                        batch_status = gr.Markdown()
                        batch_zip_out = gr.File(label="Download all (ZIP)")

                batch_generate_btn.click(
                    fn=run_batch,
                    inputs=[
                        batch_files, batch_spec_dropdown, batch_bg_color,
                        batch_zoom_slider, batch_x_slider, batch_y_slider,
                        batch_straighten_checkbox, batch_outfit_dropdown,
                    ],
                    outputs=[batch_gallery, batch_status, batch_zip_out],
                )

        gr.Markdown(
            """
            ---
            ⚠️ Photos are processed in-memory and are not stored.
            Compliance checks are automated heuristics, not a guarantee —
            verify final prints against your destination country's exact
            requirements before submission.
            """
        )

    return demo


# Concurrency is capped at 2 to protect the single CPU-basic worker from
# being driven into OOM by parallel BiRefNet inferences — each inference
# holds a 1024x1024 float32 activation stack in memory; more than a
# couple concurrent requests on 16GB shared RAM risks the crash this
# entire architecture is built to avoid. Batch requests still queue
# through the same limit — a 10-photo batch is 10 sequential dispatches
# on the caller's side (see engine.process_batch), not 10 parallel ones.
demo = build_interface()
demo.queue(max_size=20, default_concurrency_limit=2)

if __name__ == "__main__":
    logger.info("Warming up segmentation model...")
    warm_up()
    logger.info("Warm-up complete. Launching Gradio.")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(primary_hue="blue"),
    )