mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 13:30:51 +08:00
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
import numpy as np
|
|
|
|
|
|
class CTCLabelDecode():
|
|
""" Convert between text-label and text-index """
|
|
|
|
def __init__(self, character_dict_path):
|
|
super(CTCLabelDecode, self).__init__()
|
|
|
|
self.character_str = []
|
|
assert character_dict_path is not None, "character_dict_path should not be None"
|
|
|
|
if isinstance(character_dict_path, str):
|
|
with open(character_dict_path, "rb") as fin:
|
|
lines = fin.readlines()
|
|
for line in lines:
|
|
line = line.decode('utf-8').strip("\n").strip("\r\n")
|
|
self.character_str.append(line)
|
|
else:
|
|
self.character_str = character_dict_path
|
|
self.character_str.append(' ')
|
|
|
|
dict_character = self.add_special_char(self.character_str)
|
|
self.character = dict_character
|
|
|
|
self.dict = {}
|
|
for i, char in enumerate(dict_character):
|
|
self.dict[char] = i
|
|
|
|
def __call__(self, preds, label=None):
|
|
preds_idx = preds.argmax(axis=2)
|
|
preds_prob = preds.max(axis=2)
|
|
text = self.decode(preds_idx, preds_prob,
|
|
is_remove_duplicate=True)
|
|
if label is None:
|
|
return text
|
|
label = self.decode(label)
|
|
return text, label
|
|
|
|
def add_special_char(self, dict_character):
|
|
dict_character = ['blank'] + dict_character
|
|
return dict_character
|
|
|
|
def get_ignored_tokens(self):
|
|
return [0] # for ctc blank
|
|
|
|
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()
|
|
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])):
|
|
if text_index[batch_idx][idx] in ignored_tokens:
|
|
continue
|
|
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 + [1e-50])))
|
|
return result_list
|