vncgabriel/InstanceSegmentation

1# app.py2 3import gradio as gr4import torch5import numpy as np6import torch.nn.functional as F7from PIL import Image8from huggingface_hub import hf_hub_download9from inference import load_model, predict10 11# 1) Descarga automática de los pesos desde el Model Hub12model_path = hf_hub_download(13    repo_id="vncgabriel/instancia-segmentation-model",14    filename="pytorch_model.bin",15    repo_type="model",16)17 18# 2) Carga el modelo (usa GPU si está disponible)19device = torch.device("cuda" if torch.cuda.is_available() else "cpu")20model = load_model(model_path, device)21model.eval()22 23def segmentar_imagen(image: Image.Image):24    """25    Recibe una PIL Image y devuelve:26      1) overlay: imagen original con máscara semitransparente en rojo27      2) pure: máscara pura en amarillo sobre fondo negro28    """29    # Preprocesado: RGB -> numpy [H,W,3] -> tensor [1,3,H,W]30    img = image.convert("RGB")31    arr = np.array(img, dtype=np.float32) / 255.032    tensor = torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)33 34    # Padding a múltiplos de 32 (5 downsamples)35    _,_,H,W = tensor.shape36    pad_h = (32 - H % 32) % 3237    pad_w = (32 - W % 32) % 3238    tensor_p = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect")39 40    # Inferencia41    with torch.no_grad():42        mask_p = model(tensor_p)[0,0]   # [H+pad, W+pad]43 44    # Binariza y recorta al tamaño original45    mask_np = (mask_p.cpu().numpy() > 0.5).astype(np.uint8) * 25546    mask = mask_np[:H, :W]47 48    # --- Overlay en rojo ---49    overlay = img.convert("RGBA")50    mask_img = Image.fromarray(mask).convert("L")51    capa_roja = Image.new("RGBA", overlay.size, (255, 0, 0, 100))52    overlay.paste(capa_roja, mask=mask_img)53 54    # --- Máscara pura en amarillo ---55    pure = Image.new("RGB", (W, H), (0, 0, 0))56    capa_amarilla = Image.new("RGB", (W, H), (255, 255, 0))57    pure.paste(capa_amarilla, mask=mask_img)58 59    return overlay, pure60 61# 3) Interfaz Gradio62iface = gr.Interface(63    fn=segmentar_imagen,64    inputs=gr.Image(type="pil", label="Imagen de entrada"),65    outputs=[66        gr.Image(type="pil", label="Overlay en rojo"),67        gr.Image(type="pil", label="Máscara pura en amarillo"),68    ],69    title="Segmentación de Instancias (Overlay + Máscara)",70    description="Sube una imagen y obtén la segmentación de instancias: overlay rojo y máscara amarilla.",71    live=True,72)73 74if __name__ == "__main__":75    iface.launch()76