Update app.py

app.py

@@ -1,34 +1,57 @@
+# app.py
 import gradio as gr
-import numpy as np
 import torch
+import numpy as np
+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
+model_path = hf_hub_download(
+    repo_id="vncgabriel/instancia-segmentation-model",
+    filename="pytorch_model.bin",
+    repo_type="model"
+)
+
+# 2) Carrega o modelo (na CPU ou GPU)
+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)
 
-# Carrega o modelo na inicialização (GPU se disponível)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = load_model("pytorch_model.bin", device)
-
-def segment_instance(image):
-    # Converte imagem PIL em tensor normalizado [C,H,W]
-    image_array = np.array(image).astype(np.float32) / 255.0
-    image_tensor = torch.from_numpy(image_array).permute(2, 0, 1)
-    # Obtém máscara predita [1,H,W]
-    mask_tensor = predict(model, image_tensor.to(device))
-    mask = mask_tensor.squeeze(0).cpu().numpy()
-    # Limiares binários e escala para 0-255
-    mask = (mask > 0.5).astype(np.uint8) * 255
-    # Converte para 3 canais para exibição
-    mask_image = np.stack([mask]*3, axis=-1)
-    return mask_image
-
-title = "Segmentação de Instâncias"
-description = "Envie uma imagem e obtenha a segmentação de instâncias usando UNet pré-treinado."
-demo = gr.Interface(
+        # 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
+iface = gr.Interface(
     fn=segment_instance,
-    inputs=gr.Image(type="pil", label="Entrada"),
-    outputs=gr.Image(type="numpy", label="Máscara"),
-    title=title,
-    description=description
+    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."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch(server_name="0.0.0.0", server_port=7860)
+