Update inference.py

inference.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-# Definición de la arquitectura UNet (la misma utilizada en el entrenamiento).
+# Definição da arquitetura UNet (idêntica à do treinamento).
 class UNet(nn.Module):
     def __init__(self):
         super(UNet, self).__init__()
@@ -10,7 +10,7 @@ class UNet(nn.Module):
         self.encoder2 = self.conv_block(64, 128)
         self.encoder3 = self.conv_block(128, 256)
         self.encoder4 = self.conv_block(256, 512)
-        self.encoder5 = self.conv_block(512, 1024)
+        self.encoder5 = self.conv_block(512, 1024)  # Camada adicional
         self.bottleneck = self.conv_block(1024, 2048)
         self.upconv5 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2)
         self.decoder5 = self.conv_block(2048, 1024)
@@ -63,7 +63,7 @@ class UNet(nn.Module):
 
 def load_model(model_path, device='cpu'):
     """
-    Carga el modelo UNet con los pesos desde 'model_path'.
+    Carrega o modelo UNet com os pesos de 'model_path'.
     """
     model = UNet().to(device)
     model.load_state_dict(torch.load(model_path, map_location=device))
@@ -73,10 +73,10 @@ def load_model(model_path, device='cpu'):
 
 def predict(model, image_tensor):
     """
-    Realiza la predicción de la máscara de instancias para una imagen.
-    - model: modelo cargado (UNet).
+    Realiza a predição da máscara de instância para uma imagem.
+    - model: modelo carregado (UNet).
     - image_tensor: tensor FloatTensor [C,H,W] normalizado.
-    Retorna un tensor [1,H,W] con probabilidades/máscara.
+    Retorna um tensor [1,H,W] com as probabilidades/máscara.
     """
     with torch.no_grad():
         output = model(image_tensor.unsqueeze(0))