from fastai.vision.all import unet_learner, resnet34, load_learner

learn = load_learner("unet_export.pkl")

def segment_image(img):
    pred, _, _ = learn.predict(img)
    return pred

# Crear interfaz Gradio
interface = gr.Interface(
    fn=segment_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="numpy"),
    title="Segmentación Semántica con FastAI",
)

interface.launch()