alramil's picture
Fix variable interpolation in app.py
da4b751
raw
history blame contribute delete
848 Bytes
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()