Spaces:
Running
Running
import torch | |
from .ctc_postprocess import BaseRecLabelDecode | |
class CELabelDecode(BaseRecLabelDecode): | |
"""Convert between text-label and text-index.""" | |
def __init__(self, | |
character_dict_path=None, | |
use_space_char=False, | |
**kwargs): | |
super(CELabelDecode, self).__init__(character_dict_path, | |
use_space_char) | |
def __call__(self, preds, label=None, *args, **kwargs): | |
if isinstance(preds, tuple) or isinstance(preds, list): | |
preds = preds[-1] | |
if isinstance(preds, torch.Tensor): | |
preds = preds.numpy() | |
preds_idx = preds.argmax(axis=-1) | |
preds_prob = preds.max(axis=-1) | |
text = self.decode(preds_idx, preds_prob) | |
if label is None: | |
return text | |
label = self.decode(label.flatten()) | |
return text, label | |
def decode(self, text_index, text_prob=None): | |
"""convert text-index into text-label.""" | |
result_list = [] | |
batch_size = len(text_index) | |
for batch_idx in range(batch_size): | |
text = self.character[text_index[batch_idx]] | |
if text_prob is not None: | |
conf_list = text_prob[batch_idx] | |
else: | |
conf_list = 1.0 | |
result_list.append((text, conf_list)) | |
return result_list | |
def add_special_char(self, dict_character): | |
return dict_character | |