Convert between text-label and text-index for baidu warpctc
| 10 | |
| 11 | |
| 12 | class CTCLabelConverter(object): |
| 13 | """ Convert between text-label and text-index for baidu warpctc """ |
| 14 | |
| 15 | def __init__(self, flags): |
| 16 | # character (str): set of the possible characters. |
| 17 | flags = flags.Global |
| 18 | self.character_type = flags.character_type |
| 19 | self.loss_type = flags.loss_type |
| 20 | if self.character_type == 'en': |
| 21 | self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz" |
| 22 | dict_character = list(self.character_str) |
| 23 | elif self.character_type == 'ch': |
| 24 | character_dict_path = flags.character_dict_path |
| 25 | add_space = False |
| 26 | if hasattr(flags, 'use_space_char'): |
| 27 | add_space = flags.use_space_char |
| 28 | self.character_str = "" |
| 29 | with open(character_dict_path, 'rb') as fin: |
| 30 | lines = fin.readlines() |
| 31 | for line in lines: |
| 32 | line = line.decode('utf-8').strip("\n").strip("\r\n") |
| 33 | self.character_str += line |
| 34 | if add_space: |
| 35 | self.character_str += " " |
| 36 | dict_character = list(self.character_str) |
| 37 | elif self.character_type == "en_sensitive": |
| 38 | # same with ASTER setting (use 94 char). |
| 39 | self.character_str = string.printable[:-6] |
| 40 | dict_character = list(self.character_str) |
| 41 | else: |
| 42 | self.character_str = None |
| 43 | assert self.character_str is not None, \ |
| 44 | "Nonsupport type of the character: {}".format(self.character_str) |
| 45 | self.dict = {} |
| 46 | for i, char in enumerate(dict_character): |
| 47 | # NOTE: 0 is reserved for 'CTCblank' token required by CTCLoss |
| 48 | self.dict[char] = i + 1 |
| 49 | |
| 50 | self.character = ['[blank]'] + dict_character # dummy '[CTCblank]' token for CTCLoss (index 0) |
| 51 | self.char_num = len(self.character) |
| 52 | |
| 53 | def encode(self, text): |
| 54 | """convert text-label into text-index. |
| 55 | input: |
| 56 | text: text labels of each image. [batch_size] |
| 57 | output: |
| 58 | text: concatenated text index for CTCLoss. |
| 59 | [sum(text_lengths)] = [text_index_0 + text_index_1 + ... + text_index_(n - 1)] |
| 60 | length: length of each text. [batch_size] |
| 61 | """ |
| 62 | length = [len(s) for s in text] |
| 63 | # text = ''.join(text) |
| 64 | # text = [self.dict[char] for char in text] |
| 65 | d = [] |
| 66 | batch_max_length = max(length) |
| 67 | for s in text: |
| 68 | t = [self.dict[char] for char in s] |
| 69 | t.extend([0] * (batch_max_length - len(s))) |