test

app.py

@@ -9,13 +9,15 @@ def pil_to_torch(image):
     image = np.array(image)
     image = image.transpose((2, 0, 1))
     image = torch.tensor(image).float() / 255
-    return image.unsqueeze(0)
+    image = image.unsqueeze(0)
+    image = torch.nn.functional.interpolate(image.unsqueeze(0), size=(128, 128*image.shape[3]//image.shape[2]))
+    return image
 
 
 def torch_to_pil(image):
     image = image.squeeze(0).cpu().detach().numpy()
     image = image.transpose((1, 2, 0))
-    image = (image * 255).astype(np.uint8)
+    image = (np.clip(image, 0, 1) * 255).astype(np.uint8)
     image = PIL.Image.fromarray(image)
     return image
 
@@ -28,6 +30,10 @@ def image_mod(image, noise_level, denoiser):
         denoiser = dinv.models.MedianFilter()
     elif denoiser == 'BM3D':
         denoiser = dinv.models.BM3D()
+    elif denoiser == 'TV':
+        denoiser = dinv.models.TVDenoiser()
+    elif denoiser == 'TGV':
+        denoiser = dinv.models.TGVDenoiser()
     elif denoiser == 'DRUNet':
         denoiser = dinv.models.DRUNet()
     else:
@@ -44,7 +50,7 @@ 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', 'DRUNet', 'BM3D', 'MedianFilter'], value='DnCNN', label='Denoiser')
+denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'BM3D', 'MedianFilter', 'TV', 'TGV'], value='DnCNN', label='Denoiser')
 
 demo = gr.Interface(
     image_mod,