| import gradio as gr | |
| import deepinv as dinv | |
| import torch | |
| import numpy as np | |
| import PIL.Image | |
| def pil_to_torch(image, ref_size=512): | |
| image = np.array(image) | |
| image = image.transpose((2, 0, 1)) | |
| image = torch.tensor(image).float() / 255 | |
| image = image.unsqueeze(0) | |
| if ref_size == 256: | |
| size = (ref_size, ref_size) | |
| elif image.shape[2] > image.shape[3]: | |
| size = (ref_size, ref_size * image.shape[3]//image.shape[2]) | |
| else: | |
| size = (ref_size * image.shape[2]//image.shape[3], ref_size) | |
| image = torch.nn.functional.interpolate(image, size=size, mode='bilinear') | |
| return image | |
| def torch_to_pil(image): | |
| image = image.squeeze(0).cpu().detach().numpy() | |
| image = image.transpose((1, 2, 0)) | |
| image = (np.clip(image, 0, 1) * 255).astype(np.uint8) | |
| image = PIL.Image.fromarray(image) | |
| return image | |
| def image_mod(image, noise_level, denoiser): | |
| image = pil_to_torch(image, ref_size=256 if denoiser == 'DiffUNet' else 512) | |
| if denoiser == 'DnCNN': | |
| den = dinv.models.DnCNN() | |
| sigma0 = 2/255 | |
| denoiser = lambda x, sigma: den(x*sigma0/sigma)*sigma/sigma0 | |
| elif denoiser == 'MedianFilter': | |
| denoiser = dinv.models.MedianFilter(kernel_size=5) | |
| elif denoiser == 'BM3D': | |
| denoiser = dinv.models.BM3D() | |
| elif denoiser == 'TV': | |
| denoiser = dinv.models.TVDenoiser() | |
| elif denoiser == 'TGV': | |
| denoiser = dinv.models.TGVDenoiser() | |
| elif denoiser == 'Wavelets': | |
| denoiser = dinv.models.WaveletPrior() | |
| elif denoiser == 'DiffUNet': | |
| denoiser = dinv.models.DiffUNet() | |
| elif denoiser == 'DRUNet': | |
| denoiser = dinv.models.DRUNet() | |
| else: | |
| raise ValueError("Invalid denoiser") | |
| noisy = image + torch.randn_like(image) * noise_level | |
| estimated = denoiser(noisy, noise_level) | |
| return torch_to_pil(noisy), torch_to_pil(estimated) | |
| input_image = gr.Image(label='Input Image') | |
| output_images = gr.Image(label='Denoised Image') | |
| noise_image = gr.Image(label='Noisy Image') | |
| input_image_output = gr.Image(label='Input Image') | |
| noise_levels = gr.Dropdown(choices=[0.05, 0.1, 0.2, 0.3, 0.5, 1], value=0.1, label='Noise Level') | |
| denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'DiffUNet', 'BM3D', 'MedianFilter', 'TV', 'TGV', 'Wavelets'], value='DRUNet', label='Denoiser') | |
| demo = gr.Interface( | |
| image_mod, | |
| inputs=[input_image, noise_levels, denoiser], | |
| examples=[['https://upload.wikimedia.org/wikipedia/commons/b/b4/Lionel-Messi-Argentina-2022-FIFA-World-Cup_%28cropped%29.jpg', 0.1, 'DRUNet']], | |
| outputs=[noise_image, output_images], | |
| title="Image Denoising with DeepInverse", | |
| description="Denoise an image using a variety of denoisers and noise levels using the deepinverse library (https://deepinv.github.io/). We only include lightweight models like DnCNN and MedianFilter as this example is intended to be run on a CPU. We also automatically resize the input image to 512 pixels to reduce the computation time. For more advanced models, please run the code locally.", | |
| ) | |
| demo.launch() |