import gradio as gr
import torch
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
from fastai.vision.all import *

# --- Clases dummy necesarias para deserializar ---
class TargetMaskConvertTransform(ItemTransform):
    def encodes(self, x): return x

class SegmentationAlbumentationsTransform(ItemTransform):
    def __init__(self, aug=None): pass
    def encodes(self, x): return x

def get_y_fn(x): return x
def ParentSplitter(x): return x

# --- Descargar modelo ---
REPO_ID = "rugarce/model_practica3"
FILENAME = "model.pkl"

model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

learn = load_learner(model_path, cpu=True)
model = learn.model
model.eval()

# --- Inferencia simple ---
def predict(image):
    image = image.resize((640,480))
    image = np.array(image).astype(np.float32) / 255.0
    image = torch.tensor(image).permute(2,0,1).unsqueeze(0)

    with torch.no_grad():
        out = model(image)
        mask = out.argmax(dim=1).squeeze().numpy().astype(np.uint8)

    return Image.fromarray(mask * 50)

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Segmentación U-Net",
)

demo.launch()