Spaces:
Sleeping
Sleeping
| """ | |
| 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"), | |
| ) |