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