app.py

"""Image denoising runnable example.

Upload a noisy (microscopy) image; restore it. Several engines behind one contract:
  * nlm / tv / wavelet — classical denoisers (scikit-image), always runnable
  * deep               — a DnCNN trained on synthetic data at build time (Dockerfile.deep)

Contract surfaces (do not edit): /process via gr.api (also the MCP tool),
?file_url=... loader, temp-file cleanup daemon.
"""
from __future__ import annotations

import json
import os
import tempfile
import urllib.request

import gradio as gr
import numpy as np
from PIL import Image

from core import cleanup
from core.io import APP_TMP_DIR, new_tmp_path, register_protected
from core.processing import ENGINES, process, simulate_full

APP_TITLE = "Image denoising"
APP_DESCRIPTION = (
    "Upload a noisy image (PNG/JPG/TIFF). Restore it with a classical denoiser "
    "(**nlm**, **tv**, **wavelet**) or a **deep** DnCNN trained on synthetic data. "
    "A synthetic noisy microscopy image (with the clean ground truth) is baked in, so "
    "PSNR/SSIM are reported for the noisy input and the denoised result."
)
INPUT_FILE_TYPES = [".png", ".jpg", ".jpeg", ".tif", ".tiff"]

cleanup.start()


def _baked_example() -> str | None:
    p = APP_TMP_DIR / "example.tif"
    return register_protected(p) if p.exists() else None


EXAMPLE_PATH = _baked_example()


def _save_png(arr: np.ndarray, basename: str) -> str:
    out_path = new_tmp_path(basename)
    Image.fromarray(arr.astype(np.uint8)).save(out_path)
    return out_path


# ----------------- Public API (also the MCP tool) -----------------

def process_api(file_url: str, engine: str = "nlm", strength: float = 1.0) -> gr.FileData:
    """Denoise an image and return a noisy/denoised/clean summary PNG.

    Parameters
    ----------
    file_url : noisy image (PNG/JPG/TIFF) as a URL or local path
    engine : 'nlm' | 'tv' | 'wavelet' (classical) or 'deep' (DnCNN)
    strength : denoising strength (classical engines)
    """
    img, report = process(file_url, engine=engine, strength=float(strength))
    out_path = _save_png(np.asarray(img), basename="denoise.png")
    print("[process_api] " + json.dumps(report))
    return gr.FileData(path=out_path, orig_name=os.path.basename(out_path), mime_type="image/png")


def analyze_api(file_url: str, engine: str = "nlm", strength: float = 1.0) -> dict:
    """Return the denoising report (engine, PSNR/SSIM noisy vs denoised) as JSON."""
    _img, report = process(file_url, engine=engine, strength=float(strength))
    return report


# ----------------- UI -----------------

def _run_ui(file_obj, engine, strength, progress=gr.Progress(track_tqdm=True)):
    if file_obj is None:
        return None, "Upload a noisy image first."
    progress(0.3, desc=f"Denoising ({engine})…")
    r = simulate_full(file_obj, engine=engine, strength=float(strength))
    progress(1.0, desc="Done")
    return r["summary"], json.dumps(r["report"], indent=2)


with gr.Blocks(title=APP_TITLE, delete_cache=(1800, 21600)) as demo:
    gr.api(
        process_api,
        api_name="process",
        api_description=(
            "Input: a single noisy 2D IMAGE (PNG/JPG/TIFF), read as grayscale. "
            "Denoise an image (URL, file, or bytes) and return a noisy/denoised/clean "
            "summary PNG. engine: 'nlm' | 'tv' | 'wavelet' | 'deep'."
        ),
    )
    gr.api(
        analyze_api,
        api_name="analyze",
        api_description="Same input as `process` (a noisy 2D image). Return the denoising report (engine, PSNR/SSIM noisy vs denoised) as JSON.",
    )

    gr.Markdown(f"# {APP_TITLE}")
    gr.Markdown(APP_DESCRIPTION)

    last_url_state = gr.State("")

    with gr.Row():
        file_input = gr.File(file_types=INPUT_FILE_TYPES, file_count="single",
                             label="Upload a noisy image")
        with gr.Column():
            engine_input = gr.Dropdown(choices=ENGINES, value="nlm", label="Engine")
            strength_input = gr.Slider(0.2, 3.0, value=1.0, step=0.1, label="Strength (classical)")

    if EXAMPLE_PATH:
        gr.Examples(
            examples=[[EXAMPLE_PATH, "nlm", 1.0]],
            inputs=[file_input, engine_input, strength_input],
            label="Try the baked-in noisy microscopy image!",
        )

    run_btn = gr.Button("Run", variant="primary")
    with gr.Row():
        summary_out = gr.Image(label="Noisy · denoised · clean", type="numpy")
        report_box = gr.Code(label="Report", language="json")

    run_btn.click(
        fn=_run_ui,
        inputs=[file_input, engine_input, strength_input],
        outputs=[summary_out, report_box],
        show_api=False,
    )

    @demo.load(inputs=[last_url_state], outputs=[last_url_state, file_input], show_api=False)
    def _load_from_query(prev_url, request: gr.Request):
        url = (request.query_params or {}).get("file_url") or ""
        if not url or url == prev_url:
            return [gr.skip(), gr.skip()]
        suffix = os.path.splitext(url)[1] or ".png"
        fd, tmp_path = tempfile.mkstemp(suffix=suffix, dir=str(APP_TMP_DIR))
        os.close(fd)
        try:
            urllib.request.urlretrieve(url, tmp_path)
        except Exception as e:  # noqa: BLE001
            try:
                os.remove(tmp_path)
            except Exception:  # noqa: BLE001
                pass
            raise gr.Error(f"Failed to download file_url: {e}")
        return [url, gr.update(value=tmp_path)]


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=16).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        mcp_server=True,
    )