changing the appearance of the app

app.py

@@ -3,9 +3,43 @@ from PIL import Image
 import torch
 import torchvision.transforms as transforms
 import torch.nn.functional as F
+import os
+import glob
 
 from archs import create_model, resume_model
 
+# -------- Detect folders & images (assets/<folder>) --------
+IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
+
+def list_subfolders(base="assets"):
+    """Return a sorted list of immediate subfolders inside base."""
+    if not os.path.isdir(base):
+        return []
+    subs = [d for d in sorted(os.listdir(base)) if os.path.isdir(os.path.join(base, d))]
+    return subs
+
+def list_images(folder):
+    """Return full paths of images inside assets/<folder>."""
+    paths = sorted(glob.glob(os.path.join("assets", folder, "*")))
+    return [p for p in paths if p.lower().endswith(IMG_EXTS)]
+
+# -------- Folder/Gallery interactions --------
+def update_gallery(folder):
+    """Given a folder name, return the gallery items (list of image paths) and store the same list in state."""
+    files = list_images(folder)
+    return gr.update(value=files, visible=True), files 
+
+def load_from_gallery(evt: gr.SelectData, current_files):
+    """On gallery click, load the clicked image path into the input image."""
+    idx = evt.index
+    if not current_files or idx is None or idx >= len(current_files):
+        return gr.update()
+    path = current_files[idx]
+    return Image.open(path)
+
+
+# Model
+
 PATH_MODEL = './DeMoE.pt'
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_opt = {
@@ -37,19 +71,36 @@ def pad_tensor(tensor, multiple = 16):
     
     return tensor
 
-def process_img(image, task = 'auto'):
+TASK_LABELS = ["Auto", "Defocus", "Low-Light", "Global-Motion", "Synth-Global-Motion", "Local-Motion"]
+
+# Map pretty label -> internal task code used by the model
+LABEL_TO_TASK = {
+    "Auto": "auto",     
+    "Low-Light": "low_light",
+    "Global-Motion": "global_motion",  
+    "Defocus": "defocus",  
+    "Synth-Global-Motion": "synth_global_motion",
+    "Local-Motion": "local_motion",          
+}
+
+def process_img(image, task_label = 'auto'):
+    """Main inference: converts PIL -> tensor, pads, runs the model with selected task, clamps, crops, returns PIL."""
+    task_label = LABEL_TO_TASK.get(task_label, 'auto')
     tensor = pil_to_tensor(image).unsqueeze(0).to(device)
     _, _, H, W = tensor.shape
     
     tensor = pad_tensor(tensor)
 
     with torch.no_grad():
-        output = model(tensor, task)
+        output = model(tensor, task_label)
 
     output = torch.clamp(output, 0., 1.)
     output = output[:,:, :H, :W].squeeze(0)    
     return tensor_to_pil(output)
 
+
+
+
 title = 'DeMoE 🌪️​' 
 description = ''' >**Abstract**: Image deblurring, removing blurring artifacts from images, is a fundamental task in computational photography and low-level computer vision. Existing approaches focus on specialized solutions tailored to particular blur types, thus, these solutions lack generalization. This limitation in current methods implies requiring multiple models to cover several blur types, which is not practical in many real scenarios. In this paper, we introduce the first all-in-one deblurring method capable of efficiently restoring images affected by diverse blur degradations, including global motion, local motion, blur in low-light conditions, and defocus blur. We propose a mixture-of-experts (MoE) decoding module, which dynamically routes image features based on the recognized blur degradation, enabling precise and efficient restoration in an end-to-end manner. Our unified approach not only achieves performance comparable to dedicated task-specific models, but also shows promising generalization to unseen blur scenarios, particularly when leveraging appropriate expert selection.
 
@@ -66,11 +117,11 @@ Available code at [github](https://github.com/cidautai/DeMoE). More information
 <br>
 '''
 
-examples = [['examples/1POA1811.png'],
-            ['examples/12_blur.png'],
-            ['examples/0031.png'],
-            ['examples/000143.png'],
-            ['examples/blur_4.png']]
+# examples = [['examples/1POA1811.png'],
+#             ['examples/12_blur.png'],
+#             ['examples/0031.png'],
+#             ['examples/000143.png'],
+#             ['examples/blur_4.png']]
 
 css = """
     .image-frame img, .image-container img {
@@ -80,19 +131,82 @@ css = """
     }
 """
 
-demo = gr.Interface(
-    fn = process_img,
-    inputs = [
-            gr.Image(type = 'pil', label = 'input')
-    ],
-    outputs = [gr.Image(type='pil', label = 'output')],
-    title = title,
-    description = description,
-    examples = examples,
-    cache_examples=False,
-    css = css
-    
+
+# Example lists per folder under ./assets (kept simple, no helpers)
+exts = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
+
+def list_basenames(folder):
+    """Return [[basename, task_label], ...] for gr.Examples using examples_dir."""
+    paths = sorted(glob.glob(f"examples/{folder}/*"))
+    basenames = [os.path.basename(p) for p in paths if p.lower().endswith(exts)]
+    # Default task per folder (tweak as you like)
+    default_task = "auto"
+    return [[name, default_task] for name in basenames]
+
+
+examples_low_light     = list_basenames("low_light")
+examples_global_motion  = list_basenames("global_motion")
+examples_synth_global_motion = list_basenames("synth_global_motion")
+examples_local_motion      = list_basenames("local_motion")
+examples_defocus    = list_basenames("defocus")
+
+# -----------------------------
+# Gradio Blocks layout
+# -----------------------------
+with gr.Blocks(css=css, title=title) as demo:
+    gr.Markdown(f"# {title}\n\n{description}")
+
+    with gr.Row():
+        # Input image and the task selector (Radio)
+        inp_img = gr.Image(type='pil', label='input')
+        # Output image and action button
+        out_img = gr.Image(type='pil', label='output')
+    task_selector = gr.Radio(
+    choices=TASK_LABELS,
+    value="auto",
+    label="Blur type"
 )
 
+    btn = gr.Button("Restore", variant="primary")
+
+    # Connect the button to the inference function
+    btn.click(
+        fn=process_img,
+        inputs=[inp_img, task_selector],
+        outputs=[out_img]
+    )
+
+    # Examples grouped by folder (each item loads image + task automatically)
+    gr.Markdown("## Examples")
+    with gr.Row():
+        # List folders found in ./assets
+        folders = list_subfolders("examples")
+        folder_radio = gr.Radio(choices=folders, label="Examples Folders", interactive=True)
+
+        gallery = gr.Gallery(
+            label="Images from the selected folder",
+            visible=False,
+            allow_preview=True,
+            columns=6,
+            height=320,
+        )
+
+    # State holds the current file list shown in the gallery (to resolve clicks)
+    current_files_state = gr.State([])
+
+    # When changing folder -> update gallery and state
+    folder_radio.change(
+        fn=update_gallery,
+        inputs=folder_radio,
+        outputs=[gallery, current_files_state]
+    )
+
+    # When clicking a thumbnail -> load it into the input image
+    gallery.select(
+        fn=load_from_gallery,
+        inputs=[current_files_state],
+        outputs=inp_img
+    )
+
 if __name__ == '__main__':
-    demo.launch(show_error = True, ssr_mode=False)
+    demo.launch(show_error = True, server_name="0.0.0.0", server_port=7864, ssr_mode=False)

demoe_gradio.py

@@ -0,0 +1,200 @@
+import os, glob
+import gradio as gr
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import torch.nn.functional as F
+
+from archs import create_model, resume_model
+
+
+# -------- Detect folders & images (assets/<folder>) --------
+IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
+
+def list_subfolders(base="assets"):
+    """Return a sorted list of immediate subfolders inside base."""
+    if not os.path.isdir(base):
+        return []
+    subs = [d for d in sorted(os.listdir(base)) if os.path.isdir(os.path.join(base, d))]
+    return subs
+
+def list_images(folder):
+    """Return full paths of images inside assets/<folder>."""
+    paths = sorted(glob.glob(os.path.join("assets", folder, "*")))
+    return [p for p in paths if p.lower().endswith(IMG_EXTS)]
+
+
+# -------- Folder/Gallery interactions --------
+def update_gallery(folder):
+    """Given a folder name, return the gallery items (list of image paths) and store the same list in state."""
+    files = list_images(folder)
+    return gr.update(value=files, visible=True), files  # (gallery_update, state_list)
+
+def load_from_gallery(evt: gr.SelectData, current_files):
+    """On gallery click, load the clicked image path into the input image."""
+    idx = evt.index
+    if not current_files or idx is None or idx >= len(current_files):
+        return gr.update()
+    path = current_files[idx]
+    return Image.open(path)
+# -----------------------------
+# Model
+# -----------------------------
+PATH_MODEL = './DeMoE.pt'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_opt = {
+    'name': 'DeMoE',
+    'img_channels': 3,
+    'width': 32,
+    'middle_blk_num': 2,
+    'enc_blk_nums': [2, 2, 2, 2],
+    'dec_blk_nums': [2, 2, 2, 2],
+    'num_experts': 5,
+    'k_used': 1
+}
+
+pil_to_tensor = transforms.ToTensor()
+tensor_to_pil = transforms.ToPILImage()
+
+# Create and load model weights
+model = create_model(model_opt, device)
+_ = torch.load(PATH_MODEL, map_location=device, weights_only=False)  # keep compatibility with different checkpoints
+model = resume_model(model, PATH_MODEL, device)
+
+def pad_tensor(tensor, multiple=16):
+    """Pad tensor so that H and W are multiples of `multiple` (default 16)."""
+    _, _, H, W = tensor.shape
+    pad_h = (multiple - H % multiple) % multiple
+    pad_w = (multiple - W % multiple) % multiple
+    tensor = F.pad(tensor, (0, pad_w, 0, pad_h), value=0)
+    return tensor
+
+# -----------------------------
+# UI / Inference
+# -----------------------------
+title = 'DeMoE 🌪️​'
+description = ''' >**Abstract**: Image deblurring, removing blurring artifacts from images, is a fundamental task in computational photography and low-level computer vision. Existing approaches focus on specialized solutions tailored to particular blur types, thus, these solutions lack generalization. This limitation in current methods implies requiring multiple models to cover several blur types, which is not practical in many real scenarios. In this paper, we introduce the first all-in-one deblurring method capable of efficiently restoring images affected by diverse blur degradations, including global motion, local motion, blur in low-light conditions, and defocus blur. We propose a mixture-of-experts (MoE) decoding module, which dynamically routes image features based on the recognized blur degradation, enabling precise and efficient restoration in an end-to-end manner. Our unified approach not only achieves performance comparable to dedicated task-specific models, but also shows promising generalization to unseen blur scenarios, particularly when leveraging appropriate expert selection.
+[Daniel Feijoo](https://github.com/danifei), Paula Garrido-Mellado, Jaesung Rim, Álvaro García, Marcos V. Conde
+[Fundación Cidaut](https://cidaut.ai/)
+Available code at [github](https://github.com/cidautai/DeMoE). More information on the [Arxiv paper](https://arxiv.org/pdf/2508.06228). 
+> **Disclaimer:** please remember this is not a product, thus, you will notice some limitations.
+**This demo expects an image with some Low-Light degradations.**
+<br>
+'''
+
+# Visible tasks in the UI
+TASK_LABELS = ["Deblur", "Low-light", "movement", "defocus", "all"]
+
+# Map pretty label -> internal task code used by the model
+LABEL_TO_TASK = {
+    "Deblur": "global",     # change to what your model expects for general deblurring
+    "Low-light": "lowlight",
+    "movement": "local",   # if your model supports local motion blur
+    "defocus": "defocus",  # if your model supports defocus blur
+    "all": "all",          # if your model supports all types at once    
+}
+
+css = """
+    .image-frame img, .image-container img {
+        width: auto;
+        height: auto;
+        max-width: none;
+    }
+"""
+
+# Example lists per folder under ./assets (kept simple, no helpers)
+exts = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
+
+def list_basenames(folder):
+    """Return [[basename, task_label], ...] for gr.Examples using examples_dir."""
+    paths = sorted(glob.glob(f"assets/{folder}/*"))
+    basenames = [os.path.basename(p) for p in paths if p.lower().endswith(exts)]
+    # Default task per folder (tweak as you like)
+    default_task = "Low-light" if folder == "lowlight" else "Deblur"
+    return [[name, default_task] for name in basenames]
+
+examples_agentir     = list_basenames("AgentIR")
+examples_allweather  = list_basenames("allweather")
+examples_amac        = list_basenames("amac_examples")
+examples_deblur      = list_basenames("deblur")
+examples_gestures    = list_basenames("gestures")
+examples_lowlight    = list_basenames("lowlight")
+examples_monolith    = list_basenames("monolith")
+examples_superres    = list_basenames("superres")
+
+def process_img(image, task_label='auto'):
+    """Main inference: converts PIL -> tensor, pads, runs the model with selected task, clamps, crops, returns PIL."""
+    task = LABEL_TO_TASK.get(task_label, 'auto')  # default to lowlight if something unexpected arrives
+    tensor = pil_to_tensor(image).unsqueeze(0).to(device)
+    _, _, H, W = tensor.shape
+    tensor = pad_tensor(tensor)
+
+    with torch.no_grad():
+        output = model(tensor, task)
+
+    output = torch.clamp(output, 0., 1.)
+    output = output[:, :, :H, :W].squeeze(0)
+    return tensor_to_pil(output)
+
+# -----------------------------
+# Gradio Blocks layout
+# -----------------------------
+with gr.Blocks(css=css, title=title) as demo:
+    gr.Markdown(f"# {title}\n\n{description}")
+
+    with gr.Row():
+        # Input image and the task selector (Radio)
+        inp_img = gr.Image(type='pil', label='input')
+        # Output image and action button
+        out_img = gr.Image(type='pil', label='output')
+    task_selector = gr.Radio(
+    choices=TASK_LABELS,
+    value="auto",
+    label="Tipo de blur a corregir"
+)
+
+    btn = gr.Button("Corregir", variant="primary")
+
+    # Connect the button to the inference function
+    btn.click(
+        fn=process_img,
+        inputs=[inp_img, task_selector],
+        outputs=[out_img]
+    )
+
+    # Examples grouped by folder (each item loads image + task automatically)
+    gr.Markdown("## Ejemplos (assets)")
+    with gr.Row():
+        # List folders found in ./assets
+        folders = list_subfolders("assets")
+        folder_radio = gr.Radio(choices=folders, label="Carpetas en assets", interactive=True)
+
+        gallery = gr.Gallery(
+            label="Imágenes de la carpeta seleccionada",
+            visible=False,
+            allow_preview=True,
+            columns=6,
+            height=320,
+        )
+
+    # State holds the current file list shown in the gallery (to resolve clicks)
+    current_files_state = gr.State([])
+
+    # When changing folder -> update gallery and state
+    folder_radio.change(
+        fn=update_gallery,
+        inputs=folder_radio,
+        outputs=[gallery, current_files_state]
+    )
+
+    # When clicking a thumbnail -> load it into the input image
+    gallery.select(
+        fn=load_from_gallery,
+        inputs=[current_files_state],
+        outputs=inp_img
+    )
+
+
+if __name__ == '__main__':
+    # Explicit host/port and no SSR are friendly to Spaces
+    demo.launch(show_error=True, server_name="0.0.0.0", server_port=7864, ssr_mode=False)