rec_postprocess.py 23.0 KB
Newer Older
W
WenmuZhou 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
T
tink2123 已提交
15
import string
W
WenmuZhou 已提交
16 17
import paddle
from paddle.nn import functional as F
A
andyjpaddle 已提交
18
import re
W
WenmuZhou 已提交
19 20 21 22 23 24 25 26 27


class BaseRecLabelDecode(object):
    """ Convert between text-label and text-index """

    def __init__(self,
                 character_dict_path=None,
                 character_type='ch',
                 use_space_char=False):
M
MissPenguin 已提交
28
        support_character_type = [
T
tink2123 已提交
29
            'ch', 'en', 'EN_symbol', 'french', 'german', 'japan', 'korean',
T
tink2123 已提交
30 31
            'it', 'xi', 'pu', 'ru', 'ar', 'ta', 'ug', 'fa', 'ur', 'rs', 'oc',
            'rsc', 'bg', 'uk', 'be', 'te', 'ka', 'chinese_cht', 'hi', 'mr',
T
tink2123 已提交
32
            'ne', 'EN', 'latin', 'arabic', 'cyrillic', 'devanagari'
M
MissPenguin 已提交
33
        ]
W
WenmuZhou 已提交
34
        assert character_type in support_character_type, "Only {} are supported now but get {}".format(
M
MissPenguin 已提交
35
            support_character_type, character_type)
W
WenmuZhou 已提交
36

T
tink2123 已提交
37 38 39
        self.beg_str = "sos"
        self.end_str = "eos"

W
WenmuZhou 已提交
40 41 42
        if character_type == "en":
            self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz"
            dict_character = list(self.character_str)
T
tink2123 已提交
43
        elif character_type == "EN_symbol":
T
tink2123 已提交
44 45 46 47
            # same with ASTER setting (use 94 char).
            self.character_str = string.printable[:-6]
            dict_character = list(self.character_str)
        elif character_type in support_character_type:
W
WenmuZhou 已提交
48
            self.character_str = []
T
tink2123 已提交
49 50
            assert character_dict_path is not None, "character_dict_path should not be None when character_type is {}".format(
                character_type)
W
WenmuZhou 已提交
51 52 53 54
            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")
W
WenmuZhou 已提交
55
                    self.character_str.append(line)
W
WenmuZhou 已提交
56
            if use_space_char:
W
WenmuZhou 已提交
57
                self.character_str.append(" ")
W
WenmuZhou 已提交
58
            dict_character = list(self.character_str)
T
tink2123 已提交
59

W
WenmuZhou 已提交
60 61 62 63 64 65 66 67 68 69 70 71
        else:
            raise NotImplementedError
        self.character_type = character_type
        dict_character = self.add_special_char(dict_character)
        self.dict = {}
        for i, char in enumerate(dict_character):
            self.dict[char] = i
        self.character = dict_character

    def add_special_char(self, dict_character):
        return dict_character

L
littletomatodonkey 已提交
72
    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
W
WenmuZhou 已提交
73 74 75 76 77 78 79 80 81 82 83
        """ 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:
84
                    # only for predict
W
WenmuZhou 已提交
85 86 87 88 89 90 91 92 93 94
                    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)
Z
zhoujun 已提交
95
            result_list.append((text, np.mean(conf_list)))
W
WenmuZhou 已提交
96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113
        return result_list

    def get_ignored_tokens(self):
        return [0]  # for ctc blank


class CTCLabelDecode(BaseRecLabelDecode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 character_dict_path=None,
                 character_type='ch',
                 use_space_char=False,
                 **kwargs):
        super(CTCLabelDecode, self).__init__(character_dict_path,
                                             character_type, use_space_char)

    def __call__(self, preds, label=None, *args, **kwargs):
W
WenmuZhou 已提交
114 115
        if isinstance(preds, paddle.Tensor):
            preds = preds.numpy()
W
WenmuZhou 已提交
116 117
        preds_idx = preds.argmax(axis=2)
        preds_prob = preds.max(axis=2)
W
WenmuZhou 已提交
118
        text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
W
WenmuZhou 已提交
119 120
        if label is None:
            return text
L
littletomatodonkey 已提交
121
        label = self.decode(label)
W
WenmuZhou 已提交
122 123 124 125 126 127 128
        return text, label

    def add_special_char(self, dict_character):
        dict_character = ['blank'] + dict_character
        return dict_character


littletomatodonkey's avatar
littletomatodonkey 已提交
129 130 131 132 133 134 135 136 137 138
class DistillationCTCLabelDecode(CTCLabelDecode):
    """
    Convert 
    Convert between text-label and text-index
    """

    def __init__(self,
                 character_dict_path=None,
                 character_type='ch',
                 use_space_char=False,
littletomatodonkey's avatar
littletomatodonkey 已提交
139
                 model_name=["student"],
140
                 key=None,
littletomatodonkey's avatar
littletomatodonkey 已提交
141 142 143
                 **kwargs):
        super(DistillationCTCLabelDecode, self).__init__(
            character_dict_path, character_type, use_space_char)
littletomatodonkey's avatar
littletomatodonkey 已提交
144 145
        if not isinstance(model_name, list):
            model_name = [model_name]
littletomatodonkey's avatar
littletomatodonkey 已提交
146
        self.model_name = model_name
littletomatodonkey's avatar
littletomatodonkey 已提交
147

148
        self.key = key
littletomatodonkey's avatar
littletomatodonkey 已提交
149 150

    def __call__(self, preds, label=None, *args, **kwargs):
littletomatodonkey's avatar
littletomatodonkey 已提交
151 152 153 154 155 156 157
        output = dict()
        for name in self.model_name:
            pred = preds[name]
            if self.key is not None:
                pred = pred[self.key]
            output[name] = super().__call__(pred, label=label, *args, **kwargs)
        return output
littletomatodonkey's avatar
littletomatodonkey 已提交
158 159


T
Topdu 已提交
160 161 162 163 164 165 166 167 168
class NRTRLabelDecode(BaseRecLabelDecode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 character_dict_path=None,
                 character_type='EN_symbol',
                 use_space_char=True,
                 **kwargs):
        super(NRTRLabelDecode, self).__init__(character_dict_path,
A
andyjpaddle 已提交
169
                                              character_type, use_space_char)
T
Topdu 已提交
170 171 172

    def __call__(self, preds, label=None, *args, **kwargs):

T
Topdu 已提交
173 174 175 176 177 178 179 180 181 182 183 184 185
        if len(preds) == 2:
            preds_id = preds[0]
            preds_prob = preds[1]
            if isinstance(preds_id, paddle.Tensor):
                preds_id = preds_id.numpy()
            if isinstance(preds_prob, paddle.Tensor):
                preds_prob = preds_prob.numpy()
            if preds_id[0][0] == 2:
                preds_idx = preds_id[:, 1:]
                preds_prob = preds_prob[:, 1:]
            else:
                preds_idx = preds_id
            text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
T
Topdu 已提交
186 187
            if label is None:
                return text
A
andyjpaddle 已提交
188
            label = self.decode(label[:, 1:])
T
Topdu 已提交
189 190 191 192 193 194 195 196
        else:
            if isinstance(preds, paddle.Tensor):
                preds = preds.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 label is None:
                return text
A
andyjpaddle 已提交
197
            label = self.decode(label[:, 1:])
T
Topdu 已提交
198 199 200
        return text, label

    def add_special_char(self, dict_character):
A
andyjpaddle 已提交
201
        dict_character = ['blank', '<unk>', '<s>', '</s>'] + dict_character
T
Topdu 已提交
202
        return dict_character
A
andyjpaddle 已提交
203

T
Topdu 已提交
204 205 206 207 208 209 210 211
    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
        """ convert text-index into text-label. """
        result_list = []
        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])):
A
andyjpaddle 已提交
212
                if text_index[batch_idx][idx] == 3:  # end
T
Topdu 已提交
213 214
                    break
                try:
A
andyjpaddle 已提交
215 216
                    char_list.append(self.character[int(text_index[batch_idx][
                        idx])])
T
Topdu 已提交
217 218 219 220 221 222 223 224 225 226 227
                except:
                    continue
                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.lower(), np.mean(conf_list)))
        return result_list


W
WenmuZhou 已提交
228 229 230 231 232 233 234 235 236 237 238 239
class AttnLabelDecode(BaseRecLabelDecode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 character_dict_path=None,
                 character_type='ch',
                 use_space_char=False,
                 **kwargs):
        super(AttnLabelDecode, self).__init__(character_dict_path,
                                              character_type, use_space_char)

    def add_special_char(self, dict_character):
L
LDOUBLEV 已提交
240 241 242 243
        self.beg_str = "sos"
        self.end_str = "eos"
        dict_character = dict_character
        dict_character = [self.beg_str] + dict_character + [self.end_str]
W
WenmuZhou 已提交
244 245
        return dict_character

L
LDOUBLEV 已提交
246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
    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()
        [beg_idx, end_idx] = 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 int(text_index[batch_idx][idx]) == int(end_idx):
                    break
                if is_remove_duplicate:
                    # only for predict
                    if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                            batch_idx][idx]:
                        continue
A
andyjpaddle 已提交
265 266
                char_list.append(self.character[int(text_index[batch_idx][
                    idx])])
L
LDOUBLEV 已提交
267 268 269 270 271 272 273 274
                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)))
        return result_list

L
LDOUBLEV 已提交
275 276
    def __call__(self, preds, label=None, *args, **kwargs):
        """
W
WenmuZhou 已提交
277
        text = self.decode(text)
L
LDOUBLEV 已提交
278 279 280 281 282 283 284 285 286 287 288
        if label is None:
            return text
        else:
            label = self.decode(label, is_remove_duplicate=False)
            return text, label
        """
        if isinstance(preds, paddle.Tensor):
            preds = preds.numpy()

        preds_idx = preds.argmax(axis=2)
        preds_prob = preds.max(axis=2)
L
LDOUBLEV 已提交
289
        text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
L
LDOUBLEV 已提交
290 291
        if label is None:
            return text
L
LDOUBLEV 已提交
292
        label = self.decode(label, is_remove_duplicate=False)
L
LDOUBLEV 已提交
293 294
        return text, label

W
WenmuZhou 已提交
295 296 297 298 299 300 301 302 303 304 305 306 307
    def get_ignored_tokens(self):
        beg_idx = self.get_beg_end_flag_idx("beg")
        end_idx = self.get_beg_end_flag_idx("end")
        return [beg_idx, end_idx]

    def get_beg_end_flag_idx(self, beg_or_end):
        if beg_or_end == "beg":
            idx = np.array(self.dict[self.beg_str])
        elif beg_or_end == "end":
            idx = np.array(self.dict[self.end_str])
        else:
            assert False, "unsupport type %s in get_beg_end_flag_idx" \
                          % beg_or_end
M
MissPenguin 已提交
308
        return idx
T
tink2123 已提交
309 310 311 312 313 314 315 316 317 318 319 320


class SRNLabelDecode(BaseRecLabelDecode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 character_dict_path=None,
                 character_type='en',
                 use_space_char=False,
                 **kwargs):
        super(SRNLabelDecode, self).__init__(character_dict_path,
                                             character_type, use_space_char)
321
        self.max_text_length = kwargs.get('max_text_length', 25)
T
tink2123 已提交
322 323 324 325 326 327 328 329 330 331 332

    def __call__(self, preds, label=None, *args, **kwargs):
        pred = preds['predict']
        char_num = len(self.character_str) + 2
        if isinstance(pred, paddle.Tensor):
            pred = pred.numpy()
        pred = np.reshape(pred, [-1, char_num])

        preds_idx = np.argmax(pred, axis=1)
        preds_prob = np.max(pred, axis=1)

333
        preds_idx = np.reshape(preds_idx, [-1, self.max_text_length])
T
tink2123 已提交
334

335
        preds_prob = np.reshape(preds_prob, [-1, self.max_text_length])
T
tink2123 已提交
336

T
tink2123 已提交
337
        text = self.decode(preds_idx, preds_prob)
T
tink2123 已提交
338 339

        if label is None:
L
LDOUBLEV 已提交
340
            text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
T
tink2123 已提交
341
            return text
T
tink2123 已提交
342
        label = self.decode(label)
T
tink2123 已提交
343 344
        return text, label

L
LDOUBLEV 已提交
345
    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
T
tink2123 已提交
346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390
        """ 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)))
        return result_list

    def add_special_char(self, dict_character):
        dict_character = dict_character + [self.beg_str, self.end_str]
        return dict_character

    def get_ignored_tokens(self):
        beg_idx = self.get_beg_end_flag_idx("beg")
        end_idx = self.get_beg_end_flag_idx("end")
        return [beg_idx, end_idx]

    def get_beg_end_flag_idx(self, beg_or_end):
        if beg_or_end == "beg":
            idx = np.array(self.dict[self.beg_str])
        elif beg_or_end == "end":
            idx = np.array(self.dict[self.end_str])
        else:
            assert False, "unsupport type %s in get_beg_end_flag_idx" \
                          % beg_or_end
        return idx
W
WenmuZhou 已提交
391 392 393 394 395


class TableLabelDecode(object):
    """  """

A
andyjpaddle 已提交
396 397 398
    def __init__(self, character_dict_path, **kwargs):
        list_character, list_elem = self.load_char_elem_dict(
            character_dict_path)
W
WenmuZhou 已提交
399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416
        list_character = self.add_special_char(list_character)
        list_elem = self.add_special_char(list_elem)
        self.dict_character = {}
        self.dict_idx_character = {}
        for i, char in enumerate(list_character):
            self.dict_idx_character[i] = char
            self.dict_character[char] = i
        self.dict_elem = {}
        self.dict_idx_elem = {}
        for i, elem in enumerate(list_elem):
            self.dict_idx_elem[i] = elem
            self.dict_elem[elem] = i

    def load_char_elem_dict(self, character_dict_path):
        list_character = []
        list_elem = []
        with open(character_dict_path, "rb") as fin:
            lines = fin.readlines()
A
andyjpaddle 已提交
417 418
            substr = lines[0].decode('utf-8').strip("\n").strip("\r\n").split(
                "\t")
W
WenmuZhou 已提交
419 420 421
            character_num = int(substr[0])
            elem_num = int(substr[1])
            for cno in range(1, 1 + character_num):
W
WenmuZhou 已提交
422
                character = lines[cno].decode('utf-8').strip("\n").strip("\r\n")
W
WenmuZhou 已提交
423 424
                list_character.append(character)
            for eno in range(1 + character_num, 1 + character_num + elem_num):
W
WenmuZhou 已提交
425
                elem = lines[eno].decode('utf-8').strip("\n").strip("\r\n")
W
WenmuZhou 已提交
426 427 428 429 430 431 432 433 434 435 436 437
                list_elem.append(elem)
        return list_character, list_elem

    def add_special_char(self, list_character):
        self.beg_str = "sos"
        self.end_str = "eos"
        list_character = [self.beg_str] + list_character + [self.end_str]
        return list_character

    def __call__(self, preds):
        structure_probs = preds['structure_probs']
        loc_preds = preds['loc_preds']
A
andyjpaddle 已提交
438
        if isinstance(structure_probs, paddle.Tensor):
W
WenmuZhou 已提交
439
            structure_probs = structure_probs.numpy()
A
andyjpaddle 已提交
440
        if isinstance(loc_preds, paddle.Tensor):
W
WenmuZhou 已提交
441 442 443
            loc_preds = loc_preds.numpy()
        structure_idx = structure_probs.argmax(axis=2)
        structure_probs = structure_probs.max(axis=2)
A
andyjpaddle 已提交
444 445
        structure_str, structure_pos, result_score_list, result_elem_idx_list = self.decode(
            structure_idx, structure_probs, 'elem')
W
WenmuZhou 已提交
446 447 448 449 450 451 452 453 454 455 456 457 458 459
        res_html_code_list = []
        res_loc_list = []
        batch_num = len(structure_str)
        for bno in range(batch_num):
            res_loc = []
            for sno in range(len(structure_str[bno])):
                text = structure_str[bno][sno]
                if text in ['<td>', '<td']:
                    pos = structure_pos[bno][sno]
                    res_loc.append(loc_preds[bno, pos])
            res_html_code = ''.join(structure_str[bno])
            res_loc = np.array(res_loc)
            res_html_code_list.append(res_html_code)
            res_loc_list.append(res_loc)
A
andyjpaddle 已提交
460 461 462 463 464 465 466
        return {
            'res_html_code': res_html_code_list,
            'res_loc': res_loc_list,
            'res_score_list': result_score_list,
            'res_elem_idx_list': result_elem_idx_list,
            'structure_str_list': structure_str
        }
W
WenmuZhou 已提交
467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530

    def decode(self, text_index, structure_probs, char_or_elem):
        """convert text-label into text-index.
        """
        if char_or_elem == "char":
            current_dict = self.dict_idx_character
        else:
            current_dict = self.dict_idx_elem
            ignored_tokens = self.get_ignored_tokens('elem')
            beg_idx, end_idx = ignored_tokens

        result_list = []
        result_pos_list = []
        result_score_list = []
        result_elem_idx_list = []
        batch_size = len(text_index)
        for batch_idx in range(batch_size):
            char_list = []
            elem_pos_list = []
            elem_idx_list = []
            score_list = []
            for idx in range(len(text_index[batch_idx])):
                tmp_elem_idx = int(text_index[batch_idx][idx])
                if idx > 0 and tmp_elem_idx == end_idx:
                    break
                if tmp_elem_idx in ignored_tokens:
                    continue

                char_list.append(current_dict[tmp_elem_idx])
                elem_pos_list.append(idx)
                score_list.append(structure_probs[batch_idx, idx])
                elem_idx_list.append(tmp_elem_idx)
            result_list.append(char_list)
            result_pos_list.append(elem_pos_list)
            result_score_list.append(score_list)
            result_elem_idx_list.append(elem_idx_list)
        return result_list, result_pos_list, result_score_list, result_elem_idx_list

    def get_ignored_tokens(self, char_or_elem):
        beg_idx = self.get_beg_end_flag_idx("beg", char_or_elem)
        end_idx = self.get_beg_end_flag_idx("end", char_or_elem)
        return [beg_idx, end_idx]

    def get_beg_end_flag_idx(self, beg_or_end, char_or_elem):
        if char_or_elem == "char":
            if beg_or_end == "beg":
                idx = self.dict_character[self.beg_str]
            elif beg_or_end == "end":
                idx = self.dict_character[self.end_str]
            else:
                assert False, "Unsupport type %s in get_beg_end_flag_idx of char" \
                              % beg_or_end
        elif char_or_elem == "elem":
            if beg_or_end == "beg":
                idx = self.dict_elem[self.beg_str]
            elif beg_or_end == "end":
                idx = self.dict_elem[self.end_str]
            else:
                assert False, "Unsupport type %s in get_beg_end_flag_idx of elem" \
                              % beg_or_end
        else:
            assert False, "Unsupport type %s in char_or_elem" \
                          % char_or_elem
        return idx
A
andyjpaddle 已提交
531 532 533 534 535 536 537 538 539 540 541


class SARLabelDecode(BaseRecLabelDecode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 character_dict_path=None,
                 character_type='ch',
                 use_space_char=False,
                 **kwargs):
        super(SARLabelDecode, self).__init__(character_dict_path,
A
andyjpaddle 已提交
542 543
                                             character_type, use_space_char)

A
andyjpaddle 已提交
544
        self.rm_symbol = kwargs.get('rm_symbol', False)
A
andyjpaddle 已提交
545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562

    def add_special_char(self, dict_character):
        beg_end_str = "<BOS/EOS>"
        unknown_str = "<UKN>"
        padding_str = "<PAD>"
        dict_character = dict_character + [unknown_str]
        self.unknown_idx = len(dict_character) - 1
        dict_character = dict_character + [beg_end_str]
        self.start_idx = len(dict_character) - 1
        self.end_idx = len(dict_character) - 1
        dict_character = dict_character + [padding_str]
        self.padding_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()
A
andyjpaddle 已提交
563

A
andyjpaddle 已提交
564 565 566 567 568 569 570 571
        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 int(text_index[batch_idx][idx]) == int(self.end_idx):
A
andyjpaddle 已提交
572
                    if text_prob is None and idx == 0:
A
andyjpaddle 已提交
573 574 575 576 577 578 579 580 581 582 583 584 585 586 587
                        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)
588 589 590 591
            if self.rm_symbol:
                comp = re.compile('[^A-Z^a-z^0-9^\u4e00-\u9fa5]')
                text = text.lower()
                text = comp.sub('', text)
A
andyjpaddle 已提交
592 593 594 595 596 597 598 599
            result_list.append((text, np.mean(conf_list)))
        return result_list

    def __call__(self, preds, label=None, *args, **kwargs):
        if isinstance(preds, paddle.Tensor):
            preds = preds.numpy()
        preds_idx = preds.argmax(axis=2)
        preds_prob = preds.max(axis=2)
A
andyjpaddle 已提交
600

A
andyjpaddle 已提交
601 602 603 604 605 606 607 608 609
        text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)

        if label is None:
            return text
        label = self.decode(label, is_remove_duplicate=False)
        return text, label

    def get_ignored_tokens(self):
        return [self.padding_idx]