Spaces:
Running
Running
File size: 3,020 Bytes
29f689c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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]
|