test

app.py

@@ -32,7 +32,9 @@ def torch_to_pil(image):
 def image_mod(image, noise_level, denoiser):
     image = pil_to_torch(image)
     if denoiser == 'DnCNN':
-        denoiser = dinv.models.DnCNN()
+        den = dinv.models.DnCNN()
+        sigma0 = 2/255
+        denoiser = lambda x, sigma: den(x*sigma0/sigma)
     elif denoiser == 'MedianFilter':
         denoiser = dinv.models.MedianFilter(kernel_size=5)
     elif denoiser == 'BM3D':
@@ -41,6 +43,8 @@ def image_mod(image, noise_level, denoiser):
         denoiser = dinv.models.TVDenoiser()
     elif denoiser == 'TGV':
         denoiser = dinv.models.TGVDenoiser()
+    elif denoiser == 'SwinIR':
+        denoiser = dinv.models.SwinIR()
     elif denoiser == 'DRUNet':
         denoiser = dinv.models.DRUNet()
     else:
@@ -57,7 +61,7 @@ input_image_output = gr.Image(label='Input Image')
 
 noise_levels = gr.Dropdown(choices=[0.1, 0.2, 0.3, 0.5, 1], value=0.1, label='Noise Level')
 
-denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'BM3D', 'MedianFilter', 'TV', 'TGV'], value='DnCNN', label='Denoiser')
+denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'SwinIR', 'BM3D', 'MedianFilter', 'TV', 'TGV'], value='DnCNN', label='Denoiser')
 
 demo = gr.Interface(
     image_mod,