Update app.py

app.py

@@ -1,37 +1,58 @@
+# app.py
 import gradio as gr
 import torch
-import numpy as np
-from PIL import Image
 from huggingface_hub import hf_hub_download
-from model import UNet
+from inference import load_model, predict
+from PIL import Image
+import numpy as np
 
-repo_id = "vncgabriel/instance-segmentation-model"
-model_file = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
+# 1. Baixa o peso do modelo diretamente do Hugging Face Hub
+model_path = hf_hub_download(
+    repo_id="vncgabriel/instancia-segmentation-model",
+    filename="pytorch_model.pth",
+    repo_type="model",
+)
 
-model = UNet()
-state = torch.load(model_file, map_location="cpu")
-model.load_state_dict(state, strict=False)  # use strict=False to ignore mismatched keys
-model.eval()
+# 2. Carrega o modelo na CPU (ou GPU, se disponível)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+model = load_model(model_path, device=device)
+
+# 3. Função de predição para o Gradio
+def segment(image: Image.Image):
+    """
+    Recebe uma PIL Image, converte para tensor, normaliza e gera máscara.
+    Retorna imagem original sobreposta com a máscara.
+    """
+    # Converter PIL -> np.array -> tensor
+    img_arr = np.array(image).astype(np.float32) / 255.0
+    # formato HWC -> CHW
+    img_tensor = torch.from_numpy(img_arr).permute(2, 0, 1).to(device)
 
-def segment(image):
-    img = image.resize((512, 512))
-    tensor = torch.from_numpy((np.array(img)/255.0).transpose(2,0,1)).unsqueeze(0).to(device).float()
-    with torch.no_grad():
-        out = model(tensor)
-    mask = (out > 0.5).squeeze().cpu().numpy()
-    mask_img = Image.fromarray((mask*255).astype(np.uint8), mode="L").resize(image.size)
-    overlay = Image.new("RGBA", image.size, (255,0,0,80))
-    base = image.convert("RGBA")
-    base.paste(overlay, mask=mask_img)
-    return base.convert("RGB")
+    # predição
+    mask = predict(model, img_tensor).cpu().numpy()
+    # máscara binária
+    bin_mask = (mask > 0.5).astype(np.uint8) * 255
 
+    # sobrepor máscara vermelha semitransparente
+    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
+
+# 4. Define a interface Gradio
 iface = gr.Interface(
     fn=segment,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="Instance Segmentation Demo",
-    allow_flagging="never"
+    title="Segmentación de Instancias",
+    description="Sube una imagen y obtén la segmentación de instancias usando UNet preentrenado.",
+    examples=[
+        # opcional: caminhos locais ou URLs para imagens de exemplo
+        ["example1.jpg"],
+        ["example2.jpg"],
+    ],
 )
-iface.launch(server_name="0.0.0.0", server_port=7860)
+
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)