import gradio as gr
import os
import deepinv as dinv
import torch


def image_mod(image):
    denoiser = dinv.models.BM3D()
    image = torch.tensor(image).unsqueeze(0)
    image = denoiser(image)
    image = image.squeeze(0).cpu().detach().numpy()
    return image

demo = gr.Interface(
    image_mod,
    gr.Image(type="pil"),
    "image",
    flagging_options=["blurry", "incorrect", "other"],
)

demo.launch()