test

app.py

@@ -5,14 +5,15 @@ import numpy as np
 import PIL.Image
 
 
-def pil_to_torch(image):
+def pil_to_torch(image, ref_size=256):
     image = np.array(image)
     image = image.transpose((2, 0, 1))
     image = torch.tensor(image).float() / 255
     image = image.unsqueeze(0)
 
-    ref_size = 256
-    if image.shape[2] > image.shape[3]:
+    if ref_size == 128:
+        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)
@@ -30,7 +31,7 @@ def torch_to_pil(image):
 
 
 def image_mod(image, noise_level, denoiser):
-    image = pil_to_torch(image)
+    image = pil_to_torch(image, ref_size=128 if denoiser == 'SwinIR' else 256)
     if denoiser == 'DnCNN':
         den = dinv.models.DnCNN()
         sigma0 = 2/255
@@ -46,7 +47,7 @@ def image_mod(image, noise_level, denoiser):
     elif denoiser == 'Wavelets':
         denoiser = dinv.models.WaveletPrior()
     elif denoiser == 'SwinIR':
-        denoiser = dinv.models.SwinIR(img_size=256)
+        denoiser = dinv.models.SwinIR()
     elif denoiser == 'DRUNet':
         denoiser = dinv.models.DRUNet()
     else:

requirements.txt

@@ -1,4 +1,4 @@
 deepinv
 bm3d
 timm
-PyWavelets
+ptwt