Update app.py

app.py

@@ -1,37 +1,32 @@
 import gradio as gr
 import torch
 import numpy as np
-from PIL import Image, ImageDraw
+from PIL import Image
 from huggingface_hub import hf_hub_download
+from model import UNet
 
-# 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)
+state = torch.load(model_file, map_location="cpu")
+model.load_state_dict(state)
 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()
+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():
-        outputs = model(tensor)
-    mask = (outputs > 0.5).squeeze().cpu().numpy()
+        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")
 
-# Interface Gradio
 iface = gr.Interface(
     fn=segment,
     inputs=gr.Image(type="pil"),
@@ -39,5 +34,4 @@ iface = gr.Interface(
     title="Instance Segmentation Demo",
     allow_flagging="never"
 )
-
-iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch(server_name="0.0.0.0", server_port=7860)