Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| from .ctc_postprocess import BaseRecLabelDecode | |
| class SRNLabelDecode(BaseRecLabelDecode): | |
| """Convert between text-label and text-index.""" | |
| def __init__(self, | |
| character_dict_path=None, | |
| use_space_char=False, | |
| **kwargs): | |
| super(SRNLabelDecode, self).__init__(character_dict_path, | |
| use_space_char) | |
| self.max_len = 25 | |
| def add_special_char(self, dict_character): | |
| dict_character = dict_character + ['<BOS>', '<EOS>'] | |
| self.start_idx = len(dict_character) - 2 | |
| self.end_idx = len(dict_character) - 1 | |
| return dict_character | |
| def decode(self, text_index, text_prob=None, is_remove_duplicate=False): | |
| """convert text-index into text-label.""" | |
| result_list = [] | |
| ignored_tokens = self.get_ignored_tokens() | |
| # [B,25] | |
| batch_size = len(text_index) | |
| for batch_idx in range(batch_size): | |
| char_list = [] | |
| conf_list = [] | |
| for idx in range(len(text_index[batch_idx])): | |
| # print(f"text_index[{batch_idx}][{idx}]:{text_index[batch_idx][idx]}") | |
| if text_index[batch_idx][idx] in ignored_tokens: | |
| continue | |
| if int(text_index[batch_idx][idx]) == int(self.end_idx): | |
| if text_prob is None and idx == 0: | |
| continue | |
| else: | |
| break | |
| if is_remove_duplicate: | |
| # only for predict | |
| if idx > 0 and text_index[batch_idx][ | |
| idx - 1] == text_index[batch_idx][idx]: | |
| continue | |
| char_list.append(self.character[int( | |
| text_index[batch_idx][idx])]) | |
| if text_prob is not None: | |
| conf_list.append(text_prob[batch_idx][idx]) | |
| else: | |
| conf_list.append(1) | |
| text = ''.join(char_list) | |
| result_list.append((text, np.mean(conf_list).tolist())) | |
| return result_list | |
| def __call__(self, preds, batch=None, *args, **kwargs): | |
| if isinstance(preds, torch.Tensor): | |
| preds = preds.reshape([-1, self.max_len, preds.shape[-1]]) | |
| preds = preds.detach().cpu().numpy() | |
| else: | |
| preds = preds[-1] | |
| preds = preds.reshape([-1, self.max_len, | |
| preds.shape[-1]]).detach().cpu().numpy() | |
| preds_idx = preds.argmax(axis=2) | |
| preds_prob = preds.max(axis=2) | |
| text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) | |
| if batch is None: | |
| return text | |
| label = batch[1].cpu().numpy() | |
| # print(f"label.shape:{label.shape}") | |
| label = self.decode(label, is_remove_duplicate=False) | |
| return text, label | |
| def get_ignored_tokens(self): | |
| return [self.start_idx, self.end_idx] | |