add gradio app

.gitignore

@@ -6,4 +6,5 @@ temp/
 *.pdf
 *.hash
 *.npz
-sweep_results/
+sweep_results/
+.gradio

app.py

@@ -0,0 +1,201 @@
+import os
+import time
+
+import gradio as gr
+import jax
+import numpy as np
+import spaces
+import zea
+from PIL import Image
+
+from main import Config, init, run
+
+CONFIG_PATH = "configs/semantic_dps.yaml"
+SLIDER_CONFIG_PATH = "configs/slider_params.yaml"
+ASSETS_DIR = "assets"
+
+description = """
+# Semantic Diffusion Posterior Sampling for Cardiac Ultrasound Dehazing
+Select an example image below. The algorithm will dehaze the image. Note that the algorithm was heavily tuned for the DehazingEcho2025 challenge dataset, and not optimized for generalization. Therefore it is not expected to work well on any type of echocardiogram.
+
+Two parameters that are interesting to control and adjust the amount of dehazing are the "Omega (Ventricle)" and "Eta (haze prior)"
+"""
+
+
+@spaces.GPU
+def process_image(input_img, diffusion_steps, omega, omega_vent, omega_sept, eta):
+    if input_img is None:
+        raise gr.Error(
+            "No input image was provided. Please select or upload an image before running."
+        )
+
+    def _prepare_image(image):
+        resized = False
+
+        if image.mode != "L":
+            image = image.convert("L")
+
+        orig_shape = image.size[::-1]
+        h, w = diffusion_model.input_shape[:2]
+        if image.size != (w, h):
+            image = image.resize((w, h), Image.BILINEAR)
+            resized = True
+
+        image = np.array(image)
+
+        image = image.astype(np.float32)
+        image = image[None, ...]
+        return image, resized, orig_shape
+
+    try:
+        image, resized, orig_shape = _prepare_image(input_img)
+    except Exception:
+        raise gr.Error("Something went wrong with preparing the input image.")
+
+    guidance_kwargs = {
+        "omega": omega,
+        "omega_vent": omega_vent,
+        "omega_sept": omega_sept,
+        "eta": eta,
+        "smooth_l1_beta": params["guidance_kwargs"]["smooth_l1_beta"],
+    }
+
+    seed = jax.random.PRNGKey(config.seed)
+
+    try:
+        _, pred_tissue_images, *_ = run(
+            hazy_images=image,
+            diffusion_model=diffusion_model,
+            seed=seed,
+            guidance_kwargs=guidance_kwargs,
+            mask_params=params["mask_params"],
+            fixed_mask_params=params["fixed_mask_params"],
+            skeleton_params=params["skeleton_params"],
+            batch_size=1,
+            diffusion_steps=diffusion_steps,
+            initial_diffusion_step=params.get("initial_diffusion_step", 0),
+            threshold_output_quantile=params.get("threshold_output_quantile", None),
+            preserve_bottom_percent=params.get("preserve_bottom_percent", 30.0),
+            bottom_transition_width=params.get("bottom_transition_width", 10.0),
+            verbose=False,
+        )
+    except Exception:
+        raise gr.Error("The algorithm failed to process the image.")
+
+    out_img = np.squeeze(pred_tissue_images[0])
+    out_img = np.clip(out_img, 0, 255).astype(np.uint8)
+    out_pil = Image.fromarray(out_img)
+    # Resize back to original input size if needed
+    if resized and out_pil.size != (orig_shape[1], orig_shape[0]):
+        out_pil = out_pil.resize((orig_shape[1], orig_shape[0]), Image.BILINEAR)
+    # Return tuple for ImageSlider: (input, output)
+    return (input_img, out_pil)
+
+
+slider_params = Config.from_yaml(SLIDER_CONFIG_PATH)
+
+diffusion_steps_default = slider_params["diffusion_steps"]["default"]
+diffusion_steps_min = slider_params["diffusion_steps"]["min"]
+diffusion_steps_max = slider_params["diffusion_steps"]["max"]
+diffusion_steps_step = slider_params["diffusion_steps"]["step"]
+
+omega_default = slider_params["omega"]["default"]
+omega_min = slider_params["omega"]["min"]
+omega_max = slider_params["omega"]["max"]
+omega_step = slider_params["omega"]["step"]
+
+omega_vent_default = slider_params["omega_vent"]["default"]
+omega_vent_min = slider_params["omega_vent"]["min"]
+omega_vent_max = slider_params["omega_vent"]["max"]
+omega_vent_step = slider_params["omega_vent"]["step"]
+
+omega_sept_default = slider_params["omega_sept"]["default"]
+omega_sept_min = slider_params["omega_sept"]["min"]
+omega_sept_max = slider_params["omega_sept"]["max"]
+omega_sept_step = slider_params["omega_sept"]["step"]
+
+eta_default = slider_params["eta"]["default"]
+eta_min = slider_params["eta"]["min"]
+eta_max = slider_params["eta"]["max"]
+eta_step = slider_params["eta"]["step"]
+
+
+example_image_paths = [
+    os.path.join(ASSETS_DIR, f)
+    for f in os.listdir(ASSETS_DIR)
+    if f.lower().endswith(".png")
+]
+example_images = [zea.io_lib.load_image(p) for p in example_image_paths]
+examples = [[img] for img in example_images]
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(description)
+    status = gr.Markdown("Initializing model, please wait...", visible=True)
+    with gr.Row():
+        img1 = gr.Image(label="Input Image", type="pil", webcam_options=False)
+        img2 = gr.ImageSlider(label="Dehazed Image", type="pil")
+    gr.Examples(examples=examples, inputs=[img1])
+    with gr.Row():
+        diffusion_steps_slider = gr.Slider(
+            minimum=diffusion_steps_min,
+            maximum=diffusion_steps_max,
+            step=diffusion_steps_step,
+            value=diffusion_steps_default,
+            label="Diffusion Steps",
+        )
+        omega_slider = gr.Slider(
+            minimum=omega_min,
+            maximum=omega_max,
+            step=omega_step,
+            value=omega_default,
+            label="Omega (background)",
+        )
+        omega_vent_slider = gr.Slider(
+            minimum=omega_vent_min,
+            maximum=omega_vent_max,
+            step=omega_vent_step,
+            value=omega_vent_default,
+            label="Omega Ventricle",
+        )
+        omega_sept_slider = gr.Slider(
+            minimum=omega_sept_min,
+            maximum=omega_sept_max,
+            step=omega_sept_step,
+            value=omega_sept_default,
+            label="Omega Septum",
+        )
+        eta_slider = gr.Slider(
+            minimum=eta_min,
+            maximum=eta_max,
+            step=eta_step,
+            value=eta_default,
+            label="Eta (haze prior)",
+        )
+    run_btn = gr.Button("Run")
+
+    def initialize_model():
+        time.sleep(0.5)  # Let UI update
+        config = Config.from_yaml(CONFIG_PATH)
+        diffusion_model = init(config)
+        params = config.params
+        return config, diffusion_model, params
+
+    config, diffusion_model, params = initialize_model()
+    status.visible = False
+
+    run_btn.click(
+        process_image,
+        inputs=[
+            img1,
+            diffusion_steps_slider,
+            omega_slider,
+            omega_vent_slider,
+            omega_sept_slider,
+            eta_slider,
+        ],
+        outputs=[img2],
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

configs/slider_params.yaml

@@ -0,0 +1,29 @@
+diffusion_steps:
+  default: 300
+  min: 20
+  max: 500
+  step: 1
+
+omega:
+  default: 1.0
+  min: 0.5
+  max: 3.0
+  step: 0.01
+
+omega_vent:
+  default: 0.1
+  min: 0.0
+  max: 1.0
+  step: 0.01
+
+omega_sept:
+  default: 2.0
+  min: 1.0
+  max: 5.0
+  step: 0.01
+
+eta:
+  default: 0.00780
+  min: 0.0
+  max: 0.01
+  step: 0.00001

requirements.txt

@@ -0,0 +1,6 @@
+zea==0.0.4
+jax[cuda12]
+tyro
+optuna
+gradio
+spaces