import gradio as gr
import torch
import cv2
import numpy as np
from realesrgan import RealESRGAN

def upscale_image(image):
    model = RealESRGAN(torch.device('cpu'), scale=4)
    model.load_weights('weights/RealESRGAN_x4.pth', download=True)
    
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    upscaled_image = model.predict(image)
    upscaled_image = cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)
    
    return upscaled_image

iface = gr.Interface(
    fn=upscale_image,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    title="Upscaling IA con Real-ESRGAN",
    description="Carica un'immagine e migliorala con l'intelligenza artificiale!"
)

iface.launch()