app.py

from __future__ import annotations

import base64
import html
import re
import tempfile
import zipfile
from pathlib import Path

import gradio as gr

from editable_image.models import (
    BitmapExtractor,
    CleanupMode,
    DEFAULT_SNAP_ANGLES,
    Device,
    ProcessOptions,
)
from editable_image.pipeline import ProcessResult, export_result, process_path
from editable_image.svg_export import build_svg


APP_TITLE = "Editable Image"
APP_DESCRIPTION = (
    "Remove printed text from a bitmap and rebuild it as an editable SVG text layer. "
    "Shape vectorization can also promote confident diagram geometry to SVG elements."
)
ANGLE_DEFAULT = ", ".join(f"{angle:g}" for angle in DEFAULT_SNAP_ANGLES)
CSS = """
.gradio-container { max-width: 1280px !important; }
.hero { margin: 0 auto 0.6rem; max-width: 900px; text-align: center; }
.hero h1 { font-size: 2.15rem; margin-bottom: 0.25rem; }
.hero p { color: var(--body-text-color-subdued); font-size: 1rem; }
.svg-preview {
  align-items: center;
  background-color: #f4f4f5;
  background-image: linear-gradient(45deg, #e4e4e7 25%, transparent 25%),
    linear-gradient(-45deg, #e4e4e7 25%, transparent 25%),
    linear-gradient(45deg, transparent 75%, #e4e4e7 75%),
    linear-gradient(-45deg, transparent 75%, #e4e4e7 75%);
  background-position: 0 0, 0 8px, 8px -8px, -8px 0;
  background-size: 16px 16px;
  border: 1px solid var(--border-color-primary);
  border-radius: 10px;
  display: flex;
  justify-content: center;
  min-height: 420px;
  overflow: auto;
  padding: 1rem;
}
.svg-preview img { display: block; max-height: 72vh; max-width: 100%; object-fit: contain; }
"""
GRADIO_MAJOR = int(gr.__version__.split(".", maxsplit=1)[0])


def parse_snap_angles(value: str) -> list[float]:
    try:
        angles = [float(item.strip()) for item in value.split(",") if item.strip()]
    except (AttributeError, ValueError) as exc:
        raise ValueError("Snap angles must be comma-separated numbers.") from exc
    if not angles:
        raise ValueError("Enter at least one snap angle.")
    if any(angle < -180 or angle > 180 for angle in angles):
        raise ValueError("Snap angles must be between -180 and 180 degrees.")
    return angles


def safe_stem(path: Path) -> str:
    stem = re.sub(r"[^A-Za-z0-9._-]+", "-", path.stem).strip("-._")
    return stem[:100] or "image"


def file_data_uri(path: Path, media_type: str) -> str:
    payload = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:{media_type};base64,{payload}"


def preview_html(
    result: ProcessResult,
    background_path: Path,
    asset_paths: list[Path],
    embed_fonts: bool,
) -> str:
    background_uri = file_data_uri(background_path, "image/png")
    asset_by_name = {path.name: path for path in asset_paths}
    bitmap_hrefs: dict[str, tuple[str, str]] = {}
    for region in result.document.bitmap_regions:
        asset_path = asset_by_name.get(region.asset_name)
        if asset_path is not None:
            uri = file_data_uri(asset_path, "image/png")
            bitmap_hrefs[region.id] = (uri, uri)
    svg = build_svg(
        result.document,
        background_href=background_uri,
        background_fallback_href=background_uri,
        embed_fonts=embed_fonts,
        bitmap_hrefs=bitmap_hrefs,
    )
    encoded = base64.b64encode(svg.encode("utf-8")).decode("ascii")
    alt = html.escape(f"Editable reconstruction of {result.document.source_name}", quote=True)
    return (
        '<div class="svg-preview">'
        f'<img src="data:image/svg+xml;base64,{encoded}" alt="{alt}">'
        "</div>"
    )


def make_bundle(output_directory: Path, stem: str) -> Path:
    archive = output_directory.parent / f"{stem}-editable.zip"
    with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as bundle:
        for path in sorted(output_directory.rglob("*")):
            if path.is_file():
                bundle.write(path, arcname=path.relative_to(output_directory.parent))
    return archive


def convert_image(
    image_path: str | None,
    confidence: float,
    cleanup: str,
    device: str,
    max_pixels: float,
    snap_angles: str,
    snap_tolerance: float,
    snap_font_sizes: bool,
    embed_fonts: bool,
    vectorize_shapes: bool,
    bitmap_extractor: str,
) -> tuple[str, object, object, dict[str, object], str, str]:
    if not image_path:
        raise gr.Error("Upload a PNG, JPEG, or WebP image first.")

    source = Path(image_path)
    try:
        options = ProcessOptions(
            confidence=confidence,
            cleanup=CleanupMode(cleanup),
            device=Device(device),
            max_pixels=int(max_pixels),
            snap_angles=parse_snap_angles(snap_angles),
            snap_tolerance=snap_tolerance,
            snap_font_sizes=snap_font_sizes,
            vectorize_shapes=vectorize_shapes,
            bitmap_extractor=BitmapExtractor(bitmap_extractor),
        )
        result = process_path(source, options)

        job_root = Path(tempfile.mkdtemp(prefix="editable-image-"))
        stem = safe_stem(source)
        output_directory = job_root / stem
        background_path, overlay_path, asset_paths = export_result(
            result,
            output_directory,
            stem,
            embed_fonts=embed_fonts,
        )
        document_path = output_directory / f"{stem}-document.json"
        document_path.write_text(
            result.document.model_dump_json(indent=2),
            encoding="utf-8",
        )
        bundle_path = make_bundle(output_directory, stem)
        rendered = preview_html(result, background_path, asset_paths, embed_fonts)
    except Exception as exc:
        # Gradio displays this as a concise toast instead of exposing a server traceback.
        raise gr.Error(str(exc) or type(exc).__name__) from exc

    metadata: dict[str, object] = {
        "source": result.document.source_name,
        "dimensions": f"{result.image.width} × {result.image.height}",
        "recognized_text_regions": len(result.document.regions),
        "vector_shapes": len(result.document.shapes),
        "bitmap_assets": len(result.bitmap_assets),
        "cleanup": options.cleanup.value,
        "device": options.device.value,
        "document_version": result.document.version,
    }
    return (
        rendered,
        result.background,
        result.mask,
        metadata,
        str(overlay_path),
        str(bundle_path),
    )


def build_demo() -> gr.Blocks:
    blocks_options: dict[str, object] = {"title": APP_TITLE}
    if GRADIO_MAJOR < 6:
        blocks_options.update(theme=gr.themes.Soft(), css=CSS)
    with gr.Blocks(**blocks_options) as demo:
        gr.HTML(
            '<section class="hero">'
            f"<h1>{APP_TITLE}</h1><p>{APP_DESCRIPTION}</p>"
            "</section>"
        )
        with gr.Row(equal_height=False):
            with gr.Column(scale=5):
                source = gr.Image(
                    label="Source image",
                    type="filepath",
                    image_mode="RGBA",
                    sources=["upload", "clipboard"],
                    height=430,
                )
                convert = gr.Button("Make editable", variant="primary", size="lg")
                gr.Markdown(
                    "PNG, JPEG, and WebP are supported. Keep the downloaded bundle intact: "
                    "its SVG links to the sibling background and any extracted bitmap assets."
                )
            with gr.Column(scale=4):
                with gr.Accordion("OCR and orientation", open=True):
                    confidence = gr.Slider(
                        0,
                        1,
                        value=0.5,
                        step=0.01,
                        label="Minimum OCR confidence",
                    )
                    snap_angles = gr.Textbox(
                        value=ANGLE_DEFAULT,
                        label="Snap angles (degrees)",
                        info="Comma-separated targets between -180 and 180.",
                    )
                    snap_tolerance = gr.Slider(
                        0,
                        180,
                        value=6,
                        step=0.5,
                        label="Angle snap tolerance (degrees)",
                    )
                    snap_font_sizes = gr.Checkbox(
                        value=False,
                        label="Snap font sizes to detected distribution peaks",
                    )
                    embed_fonts = gr.Checkbox(
                        value=False,
                        label="Embed used WOFF2 fonts in the SVG",
                    )
                with gr.Accordion("Cleanup and vectorization", open=True):
                    cleanup = gr.Dropdown(
                        choices=[item.value for item in CleanupMode],
                        value=CleanupMode.OPENCV.value,
                        label="Cleanup method",
                        info="LaMa requires the optional dependency and installed model.",
                    )
                    vectorize_shapes = gr.Checkbox(
                        value=False,
                        label="Vectorize confident diagram shapes",
                    )
                    bitmap_extractor = gr.Dropdown(
                        choices=[item.value for item in BitmapExtractor],
                        value=BitmapExtractor.OPENCV.value,
                        label="Bitmap layer extractor",
                        info="Used during shape vectorization. SAM 2 requires its optional model.",
                    )
                with gr.Accordion("Runtime and limits", open=False):
                    device = gr.Dropdown(
                        choices=[item.value for item in Device],
                        value=Device.AUTO.value,
                        label="Compute device",
                    )
                    max_pixels = gr.Number(
                        value=40_000_000,
                        minimum=1,
                        maximum=200_000_000,
                        step=1,
                        precision=0,
                        label="Maximum decoded pixels",
                    )

        gr.Markdown("## Result")
        with gr.Tabs():
            with gr.Tab("Editable preview"):
                rendered_preview = gr.HTML(
                    '<div class="svg-preview">Upload an image and choose “Make editable”.</div>',
                    label="Reconstructed SVG",
                )
            with gr.Tab("Cleared background"):
                background = gr.Image(label="Background PNG", interactive=False)
            with gr.Tab("Removal mask"):
                mask = gr.Image(label="Applied mask", interactive=False)
            with gr.Tab("Details"):
                metadata = gr.JSON(label="Conversion details")
        with gr.Row():
            overlay = gr.File(label="Editable SVG")
            bundle = gr.File(label="Complete bundle (.zip)")

        controls = [
            source,
            confidence,
            cleanup,
            device,
            max_pixels,
            snap_angles,
            snap_tolerance,
            snap_font_sizes,
            embed_fonts,
            vectorize_shapes,
            bitmap_extractor,
        ]
        outputs = [rendered_preview, background, mask, metadata, overlay, bundle]
        convert.click(convert_image, inputs=controls, outputs=outputs, api_name="convert")

        example = Path(__file__).resolve().parent / "examples" / "scheme1.png"
        if example.exists():
            gr.Examples(
                examples=[[str(example)]],
                inputs=[source],
                label="Example",
            )
    return demo


demo = build_demo()


if __name__ == "__main__":
    launch_options: dict[str, object] = {}
    if GRADIO_MAJOR >= 6:
        launch_options.update(theme=gr.themes.Soft(), css=CSS)
    demo.queue(default_concurrency_limit=1).launch(**launch_options)