test

app.py

@@ -1,21 +1,39 @@
 import gradio as gr
-import os
 import deepinv as dinv
 import torch
+import numpy as np
 
+def pil_to_torch(image):
+    image = np.array(image)
+    image = image.transpose((2, 0, 1))
+    image = torch.tensor(image)
+    return image.unsqueeze(0)
+
+def torch_to_pil(image):
+    image = image.squeeze(0)
+    image = image.numpy()
+    image = image.transpose((1, 2, 0))
+    return image
 
 def image_mod(image):
+    image = pil_to_torch(image)
     denoiser = dinv.models.MedianFilter()
-    image = torch.tensor(image).unsqueeze(0)
-    image = denoiser(image, 0.1)
-    image = image.squeeze(0).cpu().detach().numpy()
-    return image
+    noisy = image + torch.randn_like(image) * 0.1
+    estimated = denoiser(image, 0.1)
+    return torch_to_pil(noisy), torch_to_pil(estimated)
+
+
+input_image = gr.inputs.Image()
+output_images = gr.outputs.Image(label='Denoised Image')
+noise_image = gr.outputs.Image(label='Noisy Image')
+input_image_output = gr.outputs.Image(label='Input Image')
+
 
 demo = gr.Interface(
     image_mod,
-    gr.Image(type="pil"),
-    "image",
-    flagging_options=["blurry", "incorrect", "other"],
+    inputs=input_image,
+    outputs=[noise_image, output_images],
+    title="Image Denoising with DeepInverse",
 )
 
 demo.launch()