Spaces:
Build error
Build error
File size: 2,418 Bytes
2618997 5db65e9 e64bde7 2618997 5db65e9 2618997 2ab3de8 2618997 5db65e9 2618997 5db65e9 2618997 5db65e9 2618997 5db65e9 2618997 5db65e9 e64bde7 2ab3de8 5db65e9 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | # app.py
import gradio as gr
import torch
import numpy as np
import torch.nn.functional as F
from PIL import Image
from huggingface_hub import hf_hub_download
from inference import load_model, predict
# 1) Descarga automática de los pesos desde el Model Hub
model_path = hf_hub_download(
repo_id="vncgabriel/instancia-segmentation-model",
filename="pytorch_model.bin",
repo_type="model",
)
# 2) Carga el modelo (usa GPU si está disponible)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = load_model(model_path, device)
model.eval()
def segmentar_imagen(image: Image.Image):
"""
Recibe una PIL Image y devuelve:
1) overlay: imagen original con máscara semitransparente en rojo
2) pure: máscara pura en amarillo sobre fondo negro
"""
# Preprocesado: RGB -> numpy [H,W,3] -> tensor [1,3,H,W]
img = image.convert("RGB")
arr = np.array(img, dtype=np.float32) / 255.0
tensor = torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
# Padding a múltiplos de 32 (5 downsamples)
_,_,H,W = tensor.shape
pad_h = (32 - H % 32) % 32
pad_w = (32 - W % 32) % 32
tensor_p = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect")
# Inferencia
with torch.no_grad():
mask_p = model(tensor_p)[0,0] # [H+pad, W+pad]
# Binariza y recorta al tamaño original
mask_np = (mask_p.cpu().numpy() > 0.5).astype(np.uint8) * 255
mask = mask_np[:H, :W]
# --- Overlay en rojo ---
overlay = img.convert("RGBA")
mask_img = Image.fromarray(mask).convert("L")
capa_roja = Image.new("RGBA", overlay.size, (255, 0, 0, 100))
overlay.paste(capa_roja, mask=mask_img)
# --- Máscara pura en amarillo ---
pure = Image.new("RGB", (W, H), (0, 0, 0))
capa_amarilla = Image.new("RGB", (W, H), (255, 255, 0))
pure.paste(capa_amarilla, mask=mask_img)
return overlay, pure
# 3) Interfaz Gradio
iface = gr.Interface(
fn=segmentar_imagen,
inputs=gr.Image(type="pil", label="Imagen de entrada"),
outputs=[
gr.Image(type="pil", label="Overlay en rojo"),
gr.Image(type="pil", label="Máscara pura en amarillo"),
],
title="Segmentación de Instancias (Overlay + Máscara)",
description="Sube una imagen y obtén la segmentación de instancias: overlay rojo y máscara amarilla.",
live=True,
)
if __name__ == "__main__":
iface.launch()
|