test

app.py

@@ -34,7 +34,7 @@ def image_mod(image, noise_level, denoiser):
     if denoiser == 'DnCNN':
         den = dinv.models.DnCNN()
         sigma0 = 2/255
-        denoiser = lambda x, sigma: den(x*sigma0/sigma)
+        denoiser = lambda x, sigma: den(x*sigma0/sigma)*sigma/sigma0
     elif denoiser == 'MedianFilter':
         denoiser = dinv.models.MedianFilter(kernel_size=5)
     elif denoiser == 'BM3D':
@@ -43,8 +43,10 @@ def image_mod(image, noise_level, denoiser):
         denoiser = dinv.models.TVDenoiser()
     elif denoiser == 'TGV':
         denoiser = dinv.models.TGVDenoiser()
+    elif denoiser == 'Wavelets':
+        denoiser = dinv.models.WaveletPrior()
     elif denoiser == 'SwinIR':
-        denoiser = dinv.models.SwinIR()
+        denoiser = dinv.models.SwinIR(img_size=256)
     elif denoiser == 'DRUNet':
         denoiser = dinv.models.DRUNet()
     else:
@@ -61,7 +63,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', 'SwinIR', 'BM3D', 'MedianFilter', 'TV', 'TGV'], value='DnCNN', label='Denoiser')
+denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'SwinIR', 'BM3D', 'MedianFilter', 'TV', 'TGV', 'Wavelets'], value='DnCNN', label='Denoiser')
 
 demo = gr.Interface(
     image_mod,

requirements.txt

@@ -1,2 +1,4 @@
 deepinv
-bm3d
+bm3d
+timm
+pywt