test

app.py

@@ -4,6 +4,7 @@ import torch
 import numpy as np
 import PIL.Image
 
+
 def pil_to_torch(image):
     image = np.array(image)
     image = image.transpose((2, 0, 1))
@@ -19,9 +20,14 @@ def torch_to_pil(image):
     return image
 
 
-def image_mod(image, noise_level):
+def image_mod(image, noise_level, denoiser):
     image = pil_to_torch(image)
-    denoiser = dinv.models.MedianFilter()
+    if denoiser == 'DnCNN':
+        denoiser = dinv.models.DnCNN()
+    elif denoiser == 'MedianFilter':
+        denoiser = dinv.models.MedianFilter()
+    else:
+        raise ValueError("Invalid denoiser")
     noisy = image + torch.randn_like(image) * noise_level
     estimated = denoiser(image, noise_level)
     return torch_to_pil(noisy), torch_to_pil(estimated)
@@ -34,10 +40,11 @@ input_image_output = gr.Image(label='Input Image')
 
 noise_levels = gr.Dropdown(choices=[0.1, 0.2, 0.3, 0.4, 0.5], value=0.1, label='Noise Level')
 
+denoiser = gr.Dropdown(choices=['DnCNN', 'MedianFilter'], value=0.1, label='DnCNN')
 
 demo = gr.Interface(
     image_mod,
-    inputs=[input_image, noise_levels],
+    inputs=[input_image, noise_levels, denoiser],
     outputs=[noise_image, output_images],
     title="Image Denoising with DeepInverse",
 )