| 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 | |
| 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() | |
| model = create_model(model_opt, device) | |
| checkpoints = torch.load(PATH_MODEL, map_location=device, weights_only=False) | |
| model = resume_model(model, PATH_MODEL, device) | |
| def pad_tensor(tensor, multiple = 16): | |
| '''pad the tensor to be multiple of some number''' | |
| multiple = multiple | |
| _, _, 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 | |
| def process_img(image, task = '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 = 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. | |
| [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> | |
| ''' | |
| 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 { | |
| width: auto; | |
| height: auto; | |
| max-width: none; | |
| } | |
| """ | |
| 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, | |
| css = css | |
| ) | |
| if __name__ == '__main__': | |
| demo.launch() |