convert text-label into text-index. input: text: text labels of each image. [batch_size] batch_max_length: max length of text label in the batch. 25 by default output: text : the input of attention decoder. [batch_size x (max_length+2)] +1 for [G
(self, text, batch_max_length=25)
| 116 | self.dict[char] = i |
| 117 | |
| 118 | def encode(self, text, batch_max_length=25): |
| 119 | """ convert text-label into text-index. |
| 120 | input: |
| 121 | text: text labels of each image. [batch_size] |
| 122 | batch_max_length: max length of text label in the batch. 25 by default |
| 123 | |
| 124 | output: |
| 125 | text : the input of attention decoder. [batch_size x (max_length+2)] +1 for [GO] token and +1 for [s] token. |
| 126 | text[:, 0] is [GO] token and text is padded with [GO] token after [s] token. |
| 127 | length : the length of output of attention decoder, which count [s] token also. [3, 7, ....] [batch_size] |
| 128 | """ |
| 129 | length = [len(s) + 1 for s in text] # +1 for [s] at end of sentence. |
| 130 | # batch_max_length = max(length) # this is not allowed for multi-gpu setting |
| 131 | batch_max_length += 1 |
| 132 | # additional +1 for [GO] at first step. batch_text is padded with [GO] token after [s] token. |
| 133 | batch_text = torch.LongTensor(len(text), batch_max_length + 1).fill_(0) |
| 134 | for i, t in enumerate(text): |
| 135 | text = list(t) |
| 136 | text.append('[s]') |
| 137 | text = [self.dict[char] for char in text] |
| 138 | batch_text[i][1:1 + len(text)] = torch.LongTensor(text) # batch_text[:, 0] = [GO] token |
| 139 | return (batch_text.to(device), torch.IntTensor(length).to(device)) |
| 140 | |
| 141 | def decode(self, text_index, length): |
| 142 | """ convert text-index into text-label. """ |