Update app.py

app.py

@@ -1,57 +1,75 @@
 # 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
-import traceback
 
-# 1) Baixa o peso do modelo do Hugging Face Model Hub
+# 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"
+    repo_type="model",
 )
 
-# 2) Carrega o modelo (na CPU ou GPU)
+# 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=device)
-
-def segment_instance(image: Image.Image):
-    try:
-        # 3) Preprocessamento: PIL->numpy->[0,1]->tensor [C,H,W]
-        img = image.convert("RGB")  # garante 3 canais
-        arr = np.array(img).astype(np.float32) / 255.0
-        tensor = torch.from_numpy(arr).permute(2, 0, 1).to(device)
-
-        # 4) Inferência
-        mask = predict(model, tensor).cpu().numpy()
-        mask = (mask > 0.5).astype(np.uint8) * 255  # binariza e escala
-
-        # 5) Cria overlay vermelho semi-transparente
-        overlay = img.convert("RGBA")
-        mask_img = Image.fromarray(mask).convert("L")
-        red_layer = Image.new("RGBA", overlay.size, (255, 0, 0, 100))
-        overlay.paste(red_layer, mask=mask_img)
-
-        return overlay
-
-    except Exception as e:
-        # Log completo para o console do Space
-        print(traceback.format_exc())
-        # Mostra mensagem de erro no frontend
-        return None, str(e)
-
-# 6) Monta a interface
+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=segment_instance,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Image(type="pil", label="Segmented Overlay"),
-    title="Segmentação de Instâncias",
-    description="Envie uma imagem e veja os objetos segmentados em vermelho."
+    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(server_name="0.0.0.0", server_port=7860)
-
+    iface.launch()