import gradio as gr
import numpy as np
from fastai.vision.all import load_learner, PILImage
from huggingface_hub import hf_hub_download

# Aquí sí se expande MODEL_REPO
pkl = hf_hub_download(repo_id="alramil/unet-segmentation-model", filename="export.pkl")
learn = load_learner(pkl)
learn.model.eval()

palette = np.array([
    [0,0,0],      # fondo
    [0,128,0],    # hojas
    [139,69,19],  # madera
    [128,0,128],  # poste
    [128,128,128] # racimo
], dtype=np.uint8)

def segment_image(img):
    pred,_,_ = learn.predict(PILImage.create(img))
    seg = np.array(pred)
    return palette[seg]

demo = gr.Interface(
    fn=segment_image,
    inputs=gr.Image(type="pil", label="Sube imagen"),
    outputs=gr.Image(type="numpy", label="Máscara"),
    title="Segmentación Semántica de Racimos"
)

if __name__=="__main__":
    demo.launch()