[Fix] Prob comp in vitstr and parseq for empty words (#1345)

doctr/models/recognition/parseq/pytorch.py

@@ -403,7 +403,9 @@ def __call__(
             for encoded_seq in out_idxs.cpu().numpy()
         ]
         # compute probabilties for each word up to the EOS token
-        probs = [preds_prob[i, : len(word)].clip(0, 1).mean().item() for i, word in enumerate(word_values)]
+        probs = [
+            preds_prob[i, : len(word)].clip(0, 1).mean().item() if word else 0.0 for i, word in enumerate(word_values)
+        ]
 
         return list(zip(word_values, probs))
 

doctr/models/recognition/parseq/tensorflow.py

@@ -432,7 +432,10 @@ def __call__(
         word_values = [word.decode() for word in decoded_strings_pred.numpy().tolist()]
 
         # compute probabilties for each word up to the EOS token
-        probs = [preds_prob[i, : len(word)].numpy().clip(0, 1).mean().item() for i, word in enumerate(word_values)]
+        probs = [
+            preds_prob[i, : len(word)].numpy().clip(0, 1).mean().item() if word else 0.0
+            for i, word in enumerate(word_values)
+        ]
 
         return list(zip(word_values, probs))
 

doctr/models/recognition/vitstr/pytorch.py

@@ -167,7 +167,9 @@ def __call__(
             for encoded_seq in out_idxs.cpu().numpy()
         ]
         # compute probabilties for each word up to the EOS token
-        probs = [preds_prob[i, : len(word)].clip(0, 1).mean().item() for i, word in enumerate(word_values)]
+        probs = [
+            preds_prob[i, : len(word)].clip(0, 1).mean().item() if word else 0.0 for i, word in enumerate(word_values)
+        ]
 
         return list(zip(word_values, probs))
 

doctr/models/recognition/vitstr/tensorflow.py

@@ -175,7 +175,10 @@ def __call__(
         word_values = [word.decode() for word in decoded_strings_pred.numpy().tolist()]
 
         # compute probabilties for each word up to the EOS token
-        probs = [preds_prob[i, : len(word)].numpy().clip(0, 1).mean().item() for i, word in enumerate(word_values)]
+        probs = [
+            preds_prob[i, : len(word)].numpy().clip(0, 1).mean().item() if word else 0.0
+            for i, word in enumerate(word_values)
+        ]
 
         return list(zip(word_values, probs))