Convert between text-label and text-index
| 3 | |
| 4 | |
| 5 | class CTCLabelConverter(object): |
| 6 | """ Convert between text-label and text-index """ |
| 7 | |
| 8 | def __init__(self, character): |
| 9 | # character (str): set of the possible characters. |
| 10 | dict_character = list(character) |
| 11 | |
| 12 | self.dict = {} |
| 13 | for i, char in enumerate(dict_character): |
| 14 | # NOTE: 0 is reserved for 'CTCblank' token required by CTCLoss |
| 15 | self.dict[char] = i + 1 |
| 16 | |
| 17 | self.character = ['[CTCblank]'] + dict_character # dummy '[CTCblank]' token for CTCLoss (index 0) |
| 18 | |
| 19 | def encode(self, text, batch_max_length=25): |
| 20 | """convert text-label into text-index. |
| 21 | input: |
| 22 | text: text labels of each image. [batch_size] |
| 23 | batch_max_length: max length of text label in the batch. 25 by default |
| 24 | |
| 25 | output: |
| 26 | text: text index for CTCLoss. [batch_size, batch_max_length] |
| 27 | length: length of each text. [batch_size] |
| 28 | """ |
| 29 | length = [len(s) for s in text] |
| 30 | |
| 31 | # The index used for padding (=0) would not affect the CTC loss calculation. |
| 32 | batch_text = torch.LongTensor(len(text), batch_max_length).fill_(0) |
| 33 | for i, t in enumerate(text): |
| 34 | text = list(t) |
| 35 | text = [self.dict[char] for char in text] |
| 36 | # text = [self.dict[char] if char in self.dict.keys() else 0 for char in text] |
| 37 | batch_text[i][:len(text)] = torch.LongTensor(text) |
| 38 | return (batch_text.to(device), torch.IntTensor(length).to(device)) |
| 39 | |
| 40 | def decode(self, text_index, length): |
| 41 | """ convert text-index into text-label. """ |
| 42 | texts = [] |
| 43 | for index, l in enumerate(length): |
| 44 | t = text_index[index, :] |
| 45 | |
| 46 | char_list = [] |
| 47 | for i in range(l): |
| 48 | if t[i] != 0 and (not (i > 0 and t[i - 1] == t[i])): # removing repeated characters and blank. |
| 49 | char_list.append(self.character[t[i]]) |
| 50 | text = ''.join(char_list) |
| 51 | |
| 52 | texts.append(text) |
| 53 | return texts |
| 54 | |
| 55 | |
| 56 | class CTCLabelConverterForBaiduWarpctc(object): |