File size: 848 Bytes
4515f8b
 
cfad999
4515f8b
 
da4b751
 
cfad999
 
4515f8b
 
cfad999
 
 
 
 
4515f8b
 
 
cfad999
da4b751
cfad999
4515f8b
 
 
cfad999
da4b751
 
4515f8b
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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()