Upload inference.py with huggingface_hub

inference.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-# Definição da arquitetura UNet (mesma definida no treinamento).
+# Definición de la arquitectura UNet (la misma utilizada en el entrenamiento).
 class UNet(nn.Module):
     def __init__(self):
         super(UNet, self).__init__()
@@ -23,11 +23,13 @@ class UNet(nn.Module):
         self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
         self.decoder1 = self.conv_block(128, 64)
         self.conv_last = nn.Conv2d(64, 1, kernel_size=1)
+
     def conv_block(self, in_channels, out_channels):
         return nn.Sequential(
             nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.ReLU(),
             nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.ReLU()
         )
+
     def forward(self, x):
         enc1 = self.encoder1(x)
         enc2 = self.encoder2(F.max_pool2d(enc1, 2))
@@ -35,38 +37,46 @@ class UNet(nn.Module):
         enc4 = self.encoder4(F.max_pool2d(enc3, 2))
         enc5 = self.encoder5(F.max_pool2d(enc4, 2))
         bottleneck = self.bottleneck(F.max_pool2d(enc5, 2))
+
         dec5 = self.upconv5(bottleneck)
         dec5 = torch.cat((enc5, dec5), dim=1)
         dec5 = self.decoder5(dec5)
+
         dec4 = self.upconv4(dec5)
         dec4 = torch.cat((enc4, dec4), dim=1)
         dec4 = self.decoder4(dec4)
+
         dec3 = self.upconv3(dec4)
         dec3 = torch.cat((enc3, dec3), dim=1)
         dec3 = self.decoder3(dec3)
+
         dec2 = self.upconv2(dec3)
         dec2 = torch.cat((enc2, dec2), dim=1)
         dec2 = self.decoder2(dec2)
+
         dec1 = self.upconv1(dec2)
         dec1 = torch.cat((enc1, dec1), dim=1)
         dec1 = self.decoder1(dec1)
+
         return torch.sigmoid(self.conv_last(dec1))
 
+
 def load_model(model_path, device='cpu'):
     """
-    Carrega o modelo UNet com os pesos de 'model_path'.
+    Carga el modelo UNet con los pesos desde 'model_path'.
     """
     model = UNet().to(device)
     model.load_state_dict(torch.load(model_path, map_location=device))
     model.eval()
     return model
 
+
 def predict(model, image_tensor):
     """
-    Realiza predição de máscara de instâncias para uma imagem.
-    - model: modelo carregado (UNet).
+    Realiza la predicción de la máscara de instancias para una imagen.
+    - model: modelo cargado (UNet).
     - image_tensor: tensor FloatTensor [C,H,W] normalizado.
-    Retorna tensor [1,H,W] com probabilidades/mascara.
+    Retorna un tensor [1,H,W] con probabilidades/máscara.
     """
     with torch.no_grad():
         output = model(image_tensor.unsqueeze(0))