Julián Tachella commited on
Commit
71c2965
1 Parent(s): 76c7ca0
Files changed (1) hide show
  1. app.py +7 -7
app.py CHANGED
@@ -5,13 +5,13 @@ import numpy as np
5
  import PIL.Image
6
 
7
 
8
- def pil_to_torch(image, ref_size=256):
9
  image = np.array(image)
10
  image = image.transpose((2, 0, 1))
11
  image = torch.tensor(image).float() / 255
12
  image = image.unsqueeze(0)
13
 
14
- if ref_size == 128:
15
  size = (ref_size, ref_size)
16
  elif image.shape[2] > image.shape[3]:
17
  size = (ref_size, ref_size * image.shape[3]//image.shape[2])
@@ -31,7 +31,7 @@ def torch_to_pil(image):
31
 
32
 
33
  def image_mod(image, noise_level, denoiser):
34
- image = pil_to_torch(image, ref_size=128 if denoiser == 'DiffUNet' else 256)
35
  if denoiser == 'DnCNN':
36
  den = dinv.models.DnCNN()
37
  sigma0 = 2/255
@@ -62,17 +62,17 @@ output_images = gr.Image(label='Denoised Image')
62
  noise_image = gr.Image(label='Noisy Image')
63
  input_image_output = gr.Image(label='Input Image')
64
 
65
- noise_levels = gr.Dropdown(choices=[0.1, 0.2, 0.3, 0.5, 1], value=0.1, label='Noise Level')
66
 
67
- denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'DiffUNet', 'BM3D', 'MedianFilter', 'TV', 'TGV', 'Wavelets'], value='DnCNN', label='Denoiser')
68
 
69
  demo = gr.Interface(
70
  image_mod,
71
  inputs=[input_image, noise_levels, denoiser],
72
- examples=[['https://deepinv.github.io/deepinv/_static/deepinv_logolarge.png', 0.1, 'DnCNN']],
73
  outputs=[noise_image, output_images],
74
  title="Image Denoising with DeepInverse",
75
- 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 256 pixels to reduce the computation time. For more advanced models, please run the code locally.",
76
  )
77
 
78
  demo.launch()
 
5
  import PIL.Image
6
 
7
 
8
+ def pil_to_torch(image, ref_size=512):
9
  image = np.array(image)
10
  image = image.transpose((2, 0, 1))
11
  image = torch.tensor(image).float() / 255
12
  image = image.unsqueeze(0)
13
 
14
+ if ref_size == 256:
15
  size = (ref_size, ref_size)
16
  elif image.shape[2] > image.shape[3]:
17
  size = (ref_size, ref_size * image.shape[3]//image.shape[2])
 
31
 
32
 
33
  def image_mod(image, noise_level, denoiser):
34
+ image = pil_to_torch(image, ref_size=256 if denoiser == 'DiffUNet' else 512)
35
  if denoiser == 'DnCNN':
36
  den = dinv.models.DnCNN()
37
  sigma0 = 2/255
 
62
  noise_image = gr.Image(label='Noisy Image')
63
  input_image_output = gr.Image(label='Input Image')
64
 
65
+ noise_levels = gr.Dropdown(choices=[0.05, 0.1, 0.2, 0.3, 0.5, 1], value=0.1, label='Noise Level')
66
 
67
+ denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'DiffUNet', 'BM3D', 'MedianFilter', 'TV', 'TGV', 'Wavelets'], value='DRUNet', label='Denoiser')
68
 
69
  demo = gr.Interface(
70
  image_mod,
71
  inputs=[input_image, noise_levels, denoiser],
72
+ examples=[['https://upload.wikimedia.org/wikipedia/commons/b/b4/Lionel-Messi-Argentina-2022-FIFA-World-Cup_%28cropped%29.jpg', 0.1, 'DRUNet']],
73
  outputs=[noise_image, output_images],
74
  title="Image Denoising with DeepInverse",
75
+ 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.",
76
  )
77
 
78
  demo.launch()