Create app.py

app.py

@@ -0,0 +1,43 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image, ImageDraw
+from huggingface_hub import hf_hub_download
+
+# Baixa pesos do modelo
+repo_id = "vncgabriel/instance-segmentation-model"
+model_file = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
+
+# Carrega o modelo UNet
+from __main__ import UNet
+model = UNet()
+state_dict = torch.load(model_file, map_location="cpu")
+model.load_state_dict(state_dict)
+model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# Função de inferência
+
+def segment(image: Image.Image):
+    image_resized = image.resize((512, 512))
+    tensor = torch.from_numpy((np.array(image_resized)/255.0).transpose(2,0,1)).unsqueeze(0).to(device).float()
+    with torch.no_grad():
+        outputs = model(tensor)
+    mask = (outputs > 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")
+
+# Interface Gradio
+iface = gr.Interface(
+    fn=segment,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Instance Segmentation Demo",
+    allow_flagging="never"
+)
+
+iface.launch(server_name="0.0.0.0", server_port=7860)