From f8efbce6536025f1c0fb95d3cbb4c489cb51af49 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Mon, 15 Apr 2024 21:50:20 +0200
Subject: Fix typo
---
text_recognizer/decoder/greedy_decoder.py | 24 +++++++++++++-----------
1 file changed, 13 insertions(+), 11 deletions(-)
diff --git a/text_recognizer/decoder/greedy_decoder.py b/text_recognizer/decoder/greedy_decoder.py
index 8d55a02..78864cd 100644
--- a/text_recognizer/decoder/greedy_decoder.py
+++ b/text_recognizer/decoder/greedy_decoder.py
@@ -1,8 +1,10 @@
"""Greedy decoder."""
from typing import Type
-from text_recognizer.data.tokenizer import Tokenizer
+
import torch
-from torch import nn, Tensor
+from torch import Tensor, nn
+
+from text_recognizer.data.tokenizer import Tokenizer
class GreedyDecoder:
@@ -25,29 +27,29 @@ class GreedyDecoder:
img_features = self.network.encode(x)
# Create a placeholder matrix for storing outputs from the network
- indecies = (
+ indices = (
torch.ones((bsz, self.max_output_len), dtype=torch.long, device=x.device)
* self.pad_index
)
- indecies[:, 0] = self.start_index
+ indices[:, 0] = self.start_index
for i in range(1, self.max_output_len):
- tokens = indecies[:, :i] # (B, Sy)
+ tokens = indices[:, :i] # (B, Sy)
logits = self.network.decode(tokens, img_features) # [ B, N, C ]
indecies_ = logits.argmax(dim=2) # [ B, N ]
- indecies[:, i] = indecies_[:, -1]
+ indices[:, i] = indecies_[:, -1]
# Early stopping of prediction loop if token is end or padding token.
if (
- (indecies[:, i] == self.end_index) | (indecies[:, i] == self.pad_index)
+ (indices[:, i] == self.end_index) | (indices[:, i] == self.pad_index)
).all():
break
# Set all tokens after end token to pad token.
for i in range(1, self.max_output_len):
- idx = (indecies[:, i - 1] == self.end_index) | (
- indecies[:, i - 1] == self.pad_index
+ idx = (indices[:, i - 1] == self.end_index) | (
+ indices[:, i - 1] == self.pad_index
)
- indecies[idx, i] = self.pad_index
+ indices[idx, i] = self.pad_index
- return indecies
+ return indices
--
cgit v1.2.3-70-g09d2