| 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 * |
|
|
| |
| 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 |
|
|
| |
| 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() |
|
|
| |
| 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() |
|
|
