Update app.py

app.py

@@ -1,40 +1,34 @@
-# app.py
 import gradio as gr
+import numpy as np
 import torch
-from huggingface_hub import hf_hub_download
 from inference import load_model, predict
-from PIL import Image
-import numpy as np
-
-# Baixa o peso .bin do Model Hub
-model_path = hf_hub_download(
-    repo_id="vncgabriel/instancia-segmentation-model",
-    filename="pytorch_model.bin",
-    repo_type="model",
-)
 
-# Carrega o modelo
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = load_model(model_path, device=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(image: Image.Image):
-    img_arr = np.array(image).astype(np.float32) / 255.0
-    img_tensor = torch.from_numpy(img_arr).permute(2, 0, 1).to(device)
-    mask = predict(model, img_tensor).cpu().numpy()
-    bin_mask = (mask > 0.5).astype(np.uint8) * 255
-    overlay = image.convert("RGBA")
-    mask_img = Image.fromarray(bin_mask).convert("L")
-    red = Image.new("RGBA", image.size, color=(255, 0, 0, 100))
-    overlay.paste(red, mask=mask_img)
-    return overlay
+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
 
-iface = gr.Interface(
-    fn=segment,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Segmentação de Instâncias",
-    description="Envie uma imagem e obtenha a segmentação de instâncias usando UNet pré-treinado."
+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(
+    fn=segment_instance,
+    inputs=gr.Image(type="pil", label="Entrada"),
+    outputs=gr.Image(type="numpy", label="Máscara"),
+    title=title,
+    description=description
 )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()