import gradio as gr
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()
    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,
    inputs=input_image,
    outputs=[noise_image, output_images],
    title="Image Denoising with DeepInverse",
)

demo.launch()