import cv2
import gradio as gr
import torch
from torchvision.transforms import Resize, ToTensor

from autoencoder import Autoencoder

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

model = Autoencoder()
model.load_state_dict(torch.load('model.pt', map_location=device))
model = model.eval()

resize = Resize((224))
to_tensor = ToTensor()
transforms = [to_tensor, resize]


def test(image):
    for transform in transforms:
        image = transform(image)
    image = image.unsqueeze(0)
    image = model(image).squeeze(0).permute(1,2,0).cpu().detach().numpy()
    return image


interface = gr.Interface(
    title = "OAM Autoencoder",
    description = "Select a image",
    allow_flagging="never",
    fn = test, 
    inputs = gr.Image(label = "x", type='numpy'),
    outputs = gr.Image(label = "pred"),
    examples = [
        ["img.jpg"],
    ]
)

interface.launch()