update predict function

modeling_tjmg.py

@@ -16,6 +16,7 @@ class CaptchaModel(PreTrainedModel):
         self.vocab = config.vocab
         self.output_ndigits = config.output_ndigits
         self.output_vocab_size = config.output_vocab_size
+        self.input_dim = config.input_dim
 
         self.batchnorm0 = nn.BatchNorm2d(3)
         self.conv1 = nn.Conv2d(3, 32, kernel_size=3)
@@ -48,6 +49,30 @@ class CaptchaModel(PreTrainedModel):
         for _ in range(3):
             x = self.calc_dim_img_one(x)
         return x
+    
+    def predict_captcha(self, file_path):
+        """
+        Realiza a predição do captcha para uma imagem específica.
+        """
+        # Carrega a imagem e aplica as transformações
+        transform = transforms.Compose([
+            transforms.Resize(self.input_dim),
+            transforms.ToTensor(),
+        ])
+        image = Image.open(file_path).convert('RGB')
+        image = transform(image)
+        image = image.unsqueeze(0)  # Adiciona uma dimensão para o batch
+
+        # Realiza a predição
+        with torch.no_grad():
+            logits = self.forward(image)
+
+        # Obtém a predição (índice da classe com maior probabilidade)
+        preds = torch.argmax(logits, dim=2)
+        predicted_label = "".join([self.vocab[i] for i in preds[0].tolist()])
+
+        return predicted_label
+
     def forward(self, x):
         # Passagem pela rede (observe que não usamos pipes, apenas chamadas sequenciais)
         x = self.batchnorm0(x)