from __future__ import annotations from pathlib import Path import os import tempfile import gradio as gr import numpy as np from PIL import Image from edge_aware_denoiser import ( add_gaussian_noise, as_float_image, denoise, quality_metrics, to_uint8, ) TITLE = "Image Denoiser" METHOD_NOTES = { "Perona–Malik": "Edge-preserving diffusion.", "Total Variation": "Good for strong outlines and flat areas.", "Bilateral": "Useful for light noise in photos.", } def save_png(image: np.ndarray, name: str) -> str: path = Path(tempfile.mkstemp(prefix=name, suffix=".png")[1]) Image.fromarray(to_uint8(image)).save(path) return str(path) def build_report(original, noisy, result, info, has_reference: bool) -> str: lines = [ "### Run details", f"- Method: {info.method}", f"- Estimated noise: {info.estimated_noise_sigma * 255:.2f}", ] if info.effective_kappa is not None: lines.append(f"- κ: {info.effective_kappa:.4f}") if info.iterations_completed is not None: ending = "stopped early" if info.converged else "completed" lines.append(f"- Iterations: {info.iterations_completed} ({ending})") if has_reference: noisy_scores = quality_metrics(original, noisy) result_scores = quality_metrics(original, result) lines.extend([ "", "| Metric | Noisy | Denoised |", "|---|---:|---:|", f"| RMSE | {noisy_scores['rmse']:.4f} | {result_scores['rmse']:.4f} |", f"| PSNR | {noisy_scores['psnr']:.2f} dB | {result_scores['psnr']:.2f} dB |", f"| SSIM | {noisy_scores['ssim']:.4f} | {result_scores['ssim']:.4f} |", f"| Edge similarity | {noisy_scores['edge_similarity']:.1f}% | {result_scores['edge_similarity']:.1f}% |", ]) else: edge_score = quality_metrics(noisy, result)["edge_similarity"] lines.extend(["", f"- Edge similarity to input: {edge_score:.1f}%"]) return "\n".join(lines) def process( image: np.ndarray | None, add_noise: bool, noise_sigma: float, seed: int, method: str, conductivity: str, auto_kappa: bool, manual_kappa: float, kappa_multiplier: float, iterations: int, dt: float, strength: float, spatial_radius: float, early_stop_tolerance: float, ): if image is None: raise gr.Error("Choose an image first.") original = as_float_image(image) use_noise = add_noise and noise_sigma > 0 noisy = ( add_gaussian_noise(original, noise_sigma, int(seed)) if use_noise else original.copy() ) try: result, info = denoise( noisy, method=method, conductivity=conductivity, auto_kappa=auto_kappa, manual_kappa=manual_kappa, kappa_multiplier=kappa_multiplier, iterations=int(iterations), dt=dt, strength=strength, spatial_radius=spatial_radius, early_stop_tolerance=early_stop_tolerance, ) except ValueError as error: raise gr.Error(str(error)) from error return ( to_uint8(original), to_uint8(noisy), to_uint8(result), build_report(original, noisy, result, info, use_noise), save_png(noisy, "noisy_"), save_png(result, "denoised_"), ) def update_controls(method: str, auto_kappa: bool): use_pm = method == "Perona–Malik" return ( gr.update(visible=use_pm), gr.update(visible=use_pm), gr.update(visible=use_pm and not auto_kappa), gr.update(visible=use_pm), gr.update(visible=method in {"Total Variation", "Bilateral"}), gr.update(visible=method == "Bilateral"), gr.update(value=METHOD_NOTES[method]), ) def build_app() -> gr.Blocks: with gr.Blocks(title=TITLE) as app: gr.Markdown("# Image Denoiser\nUpload an image, choose a method, and run it.") with gr.Row(): with gr.Column(scale=1, min_width=320): input_image = gr.Image( label="Image", type="numpy", image_mode="RGB", ) add_noise = gr.Checkbox(label="Add Gaussian noise", value=False) noise_sigma = gr.Slider(0, 50, value=10, step=1, label="Noise σ") seed = gr.Number(value=42, precision=0, label="Seed") method = gr.Radio( ["Perona–Malik", "Total Variation", "Bilateral"], value="Perona–Malik", label="Method", ) method_note = gr.Markdown(METHOD_NOTES["Perona–Malik"]) conductivity = gr.Dropdown( ["Exponential", "Reciprocal"], value="Exponential", label="Conductivity", ) auto_kappa = gr.Checkbox(label="Estimate κ automatically", value=True) manual_kappa = gr.Slider( 0.01, 0.30, value=0.08, step=0.005, label="κ", visible=False, ) kappa_multiplier = gr.Slider( 1.0, 5.0, value=2.8, step=0.1, label="κ multiplier", ) iterations = gr.Slider(1, 100, value=25, step=1, label="Iterations") dt = gr.Slider(0.01, 0.25, value=0.15, step=0.01, label="Time step", visible=True) strength = gr.Slider( 0.01, 0.35, value=0.10, step=0.01, label="Strength", visible=False, ) spatial_radius = gr.Slider( 1, 10, value=3, step=1, label="Spatial radius", visible=False, ) early_stop_tolerance = gr.Slider( 0.00001, 0.005, value=0.0001, step=0.00001, label="Early-stop tolerance", ) run = gr.Button("Run", variant="primary") with gr.Column(scale=2): with gr.Row(): original_view = gr.Image(label="Original", interactive=False) noisy_view = gr.Image(label="Input", interactive=False) result_view = gr.Image(label="Result", interactive=False) report = gr.Markdown("Choose an image and run denoising.") with gr.Row(): noisy_file = gr.File(label="Download input") result_file = gr.File(label="Download result") method.change( update_controls, inputs=[method, auto_kappa], outputs=[ conductivity, auto_kappa, manual_kappa, kappa_multiplier, strength, spatial_radius, method_note, ], ) auto_kappa.change( update_controls, inputs=[method, auto_kappa], outputs=[ conductivity, auto_kappa, manual_kappa, kappa_multiplier, strength, spatial_radius, method_note, ], ) run.click( process, inputs=[ input_image, add_noise, noise_sigma, seed, method, conductivity, auto_kappa, manual_kappa, kappa_multiplier, iterations, dt, strength, spatial_radius, early_stop_tolerance, ], outputs=[ original_view, noisy_view, result_view, report, noisy_file, result_file, ], ) return app if __name__ == "__main__": port = int(os.getenv("PORT", "7860")) build_app().queue(default_concurrency_limit=2, max_size=8).launch( server_name="0.0.0.0", server_port=port, inbrowser=os.getenv("LOCAL_BROWSER", "0") == "1", show_error=True, )