label_ops.py 41.8 KB
Newer Older
W
WenmuZhou 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

20
import copy
W
WenmuZhou 已提交
21
import numpy as np
T
tink2123 已提交
22
import string
L
add kie  
LDOUBLEV 已提交
23
from shapely.geometry import LineString, Point, Polygon
L
LDOUBLEV 已提交
24
import json
A
andyjpaddle 已提交
25
import copy
T
tink2123 已提交
26 27
from ppocr.utils.logging import get_logger

W
WenmuZhou 已提交
28 29 30 31 32 33 34 35 36 37 38 39

class ClsLabelEncode(object):
    def __init__(self, label_list, **kwargs):
        self.label_list = label_list

    def __call__(self, data):
        label = data['label']
        if label not in self.label_list:
            return None
        label = self.label_list.index(label)
        data['label'] = label
        return data
W
WenmuZhou 已提交
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59


class DetLabelEncode(object):
    def __init__(self, **kwargs):
        pass

    def __call__(self, data):
        label = data['label']
        label = json.loads(label)
        nBox = len(label)
        boxes, txts, txt_tags = [], [], []
        for bno in range(0, nBox):
            box = label[bno]['points']
            txt = label[bno]['transcription']
            boxes.append(box)
            txts.append(txt)
            if txt in ['*', '###']:
                txt_tags.append(True)
            else:
                txt_tags.append(False)
L
LDOUBLEV 已提交
60 61
        if len(boxes) == 0:
            return None
M
MissPenguin 已提交
62
        boxes = self.expand_points_num(boxes)
W
WenmuZhou 已提交
63 64 65 66 67 68 69 70 71 72 73 74 75
        boxes = np.array(boxes, dtype=np.float32)
        txt_tags = np.array(txt_tags, dtype=np.bool)

        data['polys'] = boxes
        data['texts'] = txts
        data['ignore_tags'] = txt_tags
        return data

    def order_points_clockwise(self, pts):
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]
        rect[2] = pts[np.argmax(s)]
L
fix  
LDOUBLEV 已提交
76 77 78 79
        tmp = np.delete(pts, (np.argmin(s), np.argmax(s)), axis=0)
        diff = np.diff(np.array(tmp), axis=1)
        rect[1] = tmp[np.argmin(diff)]
        rect[3] = tmp[np.argmax(diff)]
W
WenmuZhou 已提交
80 81
        return rect

M
MissPenguin 已提交
82 83 84 85 86 87 88 89 90 91 92
    def expand_points_num(self, boxes):
        max_points_num = 0
        for box in boxes:
            if len(box) > max_points_num:
                max_points_num = len(box)
        ex_boxes = []
        for box in boxes:
            ex_box = box + [box[-1]] * (max_points_num - len(box))
            ex_boxes.append(ex_box)
        return ex_boxes

W
WenmuZhou 已提交
93 94 95 96 97 98 99 100 101 102

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

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False):

        self.max_text_len = max_text_length
T
tink2123 已提交
103 104
        self.beg_str = "sos"
        self.end_str = "eos"
T
tink2123 已提交
105
        self.lower = False
T
tink2123 已提交
106 107 108 109 110 111

        if character_dict_path is None:
            logger = get_logger()
            logger.warning(
                "The character_dict_path is None, model can only recognize number and lower letters"
            )
W
WenmuZhou 已提交
112 113
            self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz"
            dict_character = list(self.character_str)
T
tink2123 已提交
114 115
            self.lower = True
        else:
116
            self.character_str = []
W
WenmuZhou 已提交
117 118 119 120
            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")
121
                    self.character_str.append(line)
W
WenmuZhou 已提交
122
            if use_space_char:
123
                self.character_str.append(" ")
W
WenmuZhou 已提交
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
            dict_character = list(self.character_str)
        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

    def encode(self, text):
        """convert text-label into text-index.
        input:
            text: text labels of each image. [batch_size]

        output:
            text: concatenated text index for CTCLoss.
                    [sum(text_lengths)] = [text_index_0 + text_index_1 + ... + text_index_(n - 1)]
            length: length of each text. [batch_size]
        """
W
WenmuZhou 已提交
144
        if len(text) == 0 or len(text) > self.max_text_len:
W
WenmuZhou 已提交
145
            return None
T
tink2123 已提交
146
        if self.lower:
W
WenmuZhou 已提交
147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167
            text = text.lower()
        text_list = []
        for char in text:
            if char not in self.dict:
                # logger = get_logger()
                # logger.warning('{} is not in dict'.format(char))
                continue
            text_list.append(self.dict[char])
        if len(text_list) == 0:
            return None
        return text_list


class CTCLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
168 169
        super(CTCLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
W
WenmuZhou 已提交
170 171 172 173 174 175 176 177 178

    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
        data['length'] = np.array(len(text))
        text = text + [0] * (self.max_text_len - len(text))
        data['label'] = np.array(text)
179 180 181 182 183

        label = [0] * len(self.character)
        for x in text:
            label[x] += 1
        data['label_ace'] = np.array(label)
W
WenmuZhou 已提交
184 185 186 187 188 189 190
        return data

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


J
Jethong 已提交
191
class E2ELabelEncodeTest(BaseRecLabelEncode):
J
Jethong 已提交
192 193 194 195 196
    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
197 198
        super(E2ELabelEncodeTest, self).__init__(
            max_text_length, character_dict_path, use_space_char)
J
Jethong 已提交
199 200

    def __call__(self, data):
J
Jethong 已提交
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218
        import json
        padnum = len(self.dict)
        label = data['label']
        label = json.loads(label)
        nBox = len(label)
        boxes, txts, txt_tags = [], [], []
        for bno in range(0, nBox):
            box = label[bno]['points']
            txt = label[bno]['transcription']
            boxes.append(box)
            txts.append(txt)
            if txt in ['*', '###']:
                txt_tags.append(True)
            else:
                txt_tags.append(False)
        boxes = np.array(boxes, dtype=np.float32)
        txt_tags = np.array(txt_tags, dtype=np.bool)
        data['polys'] = boxes
J
Jethong 已提交
219
        data['ignore_tags'] = txt_tags
J
Jethong 已提交
220
        temp_texts = []
J
Jethong 已提交
221
        for text in txts:
J
Jethong 已提交
222
            text = text.lower()
J
Jethong 已提交
223 224 225
            text = self.encode(text)
            if text is None:
                return None
J
Jethong 已提交
226 227
            text = text + [padnum] * (self.max_text_len - len(text)
                                      )  # use 36 to pad
J
Jethong 已提交
228 229 230 231 232
            temp_texts.append(text)
        data['texts'] = np.array(temp_texts)
        return data


J
Jethong 已提交
233
class E2ELabelEncodeTrain(object):
J
Jethong 已提交
234 235
    def __init__(self, **kwargs):
        pass
J
Jethong 已提交
236 237

    def __call__(self, data):
J
Jethong 已提交
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256
        import json
        label = data['label']
        label = json.loads(label)
        nBox = len(label)
        boxes, txts, txt_tags = [], [], []
        for bno in range(0, nBox):
            box = label[bno]['points']
            txt = label[bno]['transcription']
            boxes.append(box)
            txts.append(txt)
            if txt in ['*', '###']:
                txt_tags.append(True)
            else:
                txt_tags.append(False)
        boxes = np.array(boxes, dtype=np.float32)
        txt_tags = np.array(txt_tags, dtype=np.bool)

        data['polys'] = boxes
        data['texts'] = txts
J
Jethong 已提交
257
        data['ignore_tags'] = txt_tags
J
Jethong 已提交
258 259 260
        return data


L
add kie  
LDOUBLEV 已提交
261
class KieLabelEncode(object):
262 263 264 265 266 267
    def __init__(self,
                 character_dict_path,
                 class_path,
                 norm=10,
                 directed=False,
                 **kwargs):
L
add kie  
LDOUBLEV 已提交
268 269
        super(KieLabelEncode, self).__init__()
        self.dict = dict({'': 0})
270
        self.label2classid_map = dict()
L
fix win  
LDOUBLEV 已提交
271
        with open(character_dict_path, 'r', encoding='utf-8') as fr:
L
add kie  
LDOUBLEV 已提交
272 273 274 275 276
            idx = 1
            for line in fr:
                char = line.strip()
                self.dict[char] = idx
                idx += 1
277 278 279 280 281
        with open(class_path, "r") as fin:
            lines = fin.readlines()
            for idx, line in enumerate(lines):
                line = line.strip("\n")
                self.label2classid_map[line] = idx
L
add kie  
LDOUBLEV 已提交
282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299
        self.norm = norm
        self.directed = directed

    def compute_relation(self, boxes):
        """Compute relation between every two boxes."""
        x1s, y1s = boxes[:, 0:1], boxes[:, 1:2]
        x2s, y2s = boxes[:, 4:5], boxes[:, 5:6]
        ws, hs = x2s - x1s + 1, np.maximum(y2s - y1s + 1, 1)
        dxs = (x1s[:, 0][None] - x1s) / self.norm
        dys = (y1s[:, 0][None] - y1s) / self.norm
        xhhs, xwhs = hs[:, 0][None] / hs, ws[:, 0][None] / hs
        whs = ws / hs + np.zeros_like(xhhs)
        relations = np.stack([dxs, dys, whs, xhhs, xwhs], -1)
        bboxes = np.concatenate([x1s, y1s, x2s, y2s], -1).astype(np.float32)
        return relations, bboxes

    def pad_text_indices(self, text_inds):
        """Pad text index to same length."""
L
debug  
LDOUBLEV 已提交
300
        max_len = 300
L
add kie  
LDOUBLEV 已提交
301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325
        recoder_len = max([len(text_ind) for text_ind in text_inds])
        padded_text_inds = -np.ones((len(text_inds), max_len), np.int32)
        for idx, text_ind in enumerate(text_inds):
            padded_text_inds[idx, :len(text_ind)] = np.array(text_ind)
        return padded_text_inds, recoder_len

    def list_to_numpy(self, ann_infos):
        """Convert bboxes, relations, texts and labels to ndarray."""
        boxes, text_inds = ann_infos['points'], ann_infos['text_inds']
        boxes = np.array(boxes, np.int32)
        relations, bboxes = self.compute_relation(boxes)

        labels = ann_infos.get('labels', None)
        if labels is not None:
            labels = np.array(labels, np.int32)
            edges = ann_infos.get('edges', None)
            if edges is not None:
                labels = labels[:, None]
                edges = np.array(edges)
                edges = (edges[:, None] == edges[None, :]).astype(np.int32)
                if self.directed:
                    edges = (edges & labels == 1).astype(np.int32)
                np.fill_diagonal(edges, -1)
                labels = np.concatenate([labels, edges], -1)
        padded_text_inds, recoder_len = self.pad_text_indices(text_inds)
L
debug  
LDOUBLEV 已提交
326
        max_num = 300
L
add kie  
LDOUBLEV 已提交
327 328
        temp_bboxes = np.zeros([max_num, 4])
        h, _ = bboxes.shape
那珈落's avatar
那珈落 已提交
329
        temp_bboxes[:h, :] = bboxes
L
add kie  
LDOUBLEV 已提交
330 331 332 333

        temp_relations = np.zeros([max_num, max_num, 5])
        temp_relations[:h, :h, :] = relations

L
debug  
LDOUBLEV 已提交
334
        temp_padded_text_inds = np.zeros([max_num, max_num])
L
add kie  
LDOUBLEV 已提交
335 336
        temp_padded_text_inds[:h, :] = padded_text_inds

L
debug  
LDOUBLEV 已提交
337
        temp_labels = np.zeros([max_num, max_num])
L
add kie  
LDOUBLEV 已提交
338 339 340 341 342 343 344 345 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 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420
        temp_labels[:h, :h + 1] = labels

        tag = np.array([h, recoder_len])
        return dict(
            image=ann_infos['image'],
            points=temp_bboxes,
            relations=temp_relations,
            texts=temp_padded_text_inds,
            labels=temp_labels,
            tag=tag)

    def convert_canonical(self, points_x, points_y):

        assert len(points_x) == 4
        assert len(points_y) == 4

        points = [Point(points_x[i], points_y[i]) for i in range(4)]

        polygon = Polygon([(p.x, p.y) for p in points])
        min_x, min_y, _, _ = polygon.bounds
        points_to_lefttop = [
            LineString([points[i], Point(min_x, min_y)]) for i in range(4)
        ]
        distances = np.array([line.length for line in points_to_lefttop])
        sort_dist_idx = np.argsort(distances)
        lefttop_idx = sort_dist_idx[0]

        if lefttop_idx == 0:
            point_orders = [0, 1, 2, 3]
        elif lefttop_idx == 1:
            point_orders = [1, 2, 3, 0]
        elif lefttop_idx == 2:
            point_orders = [2, 3, 0, 1]
        else:
            point_orders = [3, 0, 1, 2]

        sorted_points_x = [points_x[i] for i in point_orders]
        sorted_points_y = [points_y[j] for j in point_orders]

        return sorted_points_x, sorted_points_y

    def sort_vertex(self, points_x, points_y):

        assert len(points_x) == 4
        assert len(points_y) == 4

        x = np.array(points_x)
        y = np.array(points_y)
        center_x = np.sum(x) * 0.25
        center_y = np.sum(y) * 0.25

        x_arr = np.array(x - center_x)
        y_arr = np.array(y - center_y)

        angle = np.arctan2(y_arr, x_arr) * 180.0 / np.pi
        sort_idx = np.argsort(angle)

        sorted_points_x, sorted_points_y = [], []
        for i in range(4):
            sorted_points_x.append(points_x[sort_idx[i]])
            sorted_points_y.append(points_y[sort_idx[i]])

        return self.convert_canonical(sorted_points_x, sorted_points_y)

    def __call__(self, data):
        import json
        label = data['label']
        annotations = json.loads(label)
        boxes, texts, text_inds, labels, edges = [], [], [], [], []
        for ann in annotations:
            box = ann['points']
            x_list = [box[i][0] for i in range(4)]
            y_list = [box[i][1] for i in range(4)]
            sorted_x_list, sorted_y_list = self.sort_vertex(x_list, y_list)
            sorted_box = []
            for x, y in zip(sorted_x_list, sorted_y_list):
                sorted_box.append(x)
                sorted_box.append(y)
            boxes.append(sorted_box)
            text = ann['transcription']
            texts.append(ann['transcription'])
            text_ind = [self.dict[c] for c in text if c in self.dict]
            text_inds.append(text_ind)
L
fix  
LDOUBLEV 已提交
421
            if 'label' in ann.keys():
422
                labels.append(self.label2classid_map[ann['label']])
L
fix  
LDOUBLEV 已提交
423 424
            elif 'key_cls' in ann.keys():
                labels.append(ann['key_cls'])
L
fix  
LDOUBLEV 已提交
425
            else:
文幕地方's avatar
文幕地方 已提交
426 427 428
                raise ValueError(
                    "Cannot found 'key_cls' in ann.keys(), please check your training annotation."
                )
L
add kie  
LDOUBLEV 已提交
429 430 431 432 433 434 435 436 437 438 439 440
            edges.append(ann.get('edge', 0))
        ann_infos = dict(
            image=data['image'],
            points=boxes,
            texts=texts,
            text_inds=text_inds,
            edges=edges,
            labels=labels)

        return self.list_to_numpy(ann_infos)


W
WenmuZhou 已提交
441 442 443 444 445 446 447 448
class AttnLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
449 450
        super(AttnLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
W
WenmuZhou 已提交
451 452

    def add_special_char(self, dict_character):
L
LDOUBLEV 已提交
453 454 455
        self.beg_str = "sos"
        self.end_str = "eos"
        dict_character = [self.beg_str] + dict_character + [self.end_str]
W
WenmuZhou 已提交
456 457
        return dict_character

L
LDOUBLEV 已提交
458 459
    def __call__(self, data):
        text = data['label']
W
WenmuZhou 已提交
460
        text = self.encode(text)
L
LDOUBLEV 已提交
461 462
        if text is None:
            return None
L
LDOUBLEV 已提交
463
        if len(text) >= self.max_text_len:
L
LDOUBLEV 已提交
464 465 466
            return None
        data['length'] = np.array(len(text))
        text = [0] + text + [len(self.character) - 1] + [0] * (self.max_text_len
T
tink2123 已提交
467
                                                               - len(text) - 2)
L
LDOUBLEV 已提交
468 469 470 471 472 473 474
        data['label'] = np.array(text)
        return data

    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]
W
WenmuZhou 已提交
475 476 477 478 479 480 481 482 483 484

    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
T
tink2123 已提交
485 486


T
tink2123 已提交
487 488 489 490 491 492 493 494
class SEEDLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
495 496
        super(SEEDLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
T
tink2123 已提交
497 498

    def add_special_char(self, dict_character):
T
tink2123 已提交
499
        self.padding = "padding"
T
tink2123 已提交
500
        self.end_str = "eos"
T
tink2123 已提交
501 502 503 504
        self.unknown = "unknown"
        dict_character = dict_character + [
            self.end_str, self.padding, self.unknown
        ]
T
tink2123 已提交
505 506 507 508 509 510 511 512 513
        return dict_character

    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
        if len(text) >= self.max_text_len:
            return None
T
rm anno  
tink2123 已提交
514
        data['length'] = np.array(len(text)) + 1  # conclude eos
T
tink2123 已提交
515 516
        text = text + [len(self.character) - 3] + [len(self.character) - 2] * (
            self.max_text_len - len(text) - 1)
T
tink2123 已提交
517 518 519 520
        data['label'] = np.array(text)
        return data


T
tink2123 已提交
521 522 523 524 525 526 527 528
class SRNLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length=25,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
529 530
        super(SRNLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
T
tink2123 已提交
531 532 533 534 535 536 537 538

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

    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
T
tink2123 已提交
539
        char_num = len(self.character)
T
tink2123 已提交
540 541 542 543 544
        if text is None:
            return None
        if len(text) > self.max_text_len:
            return None
        data['length'] = np.array(len(text))
T
tink2123 已提交
545
        text = text + [char_num - 1] * (self.max_text_len - len(text))
T
tink2123 已提交
546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562
        data['label'] = np.array(text)
        return data

    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
M
MissPenguin 已提交
563

L
LDOUBLEV 已提交
564

文幕地方's avatar
文幕地方 已提交
565
class TableLabelEncode(AttnLabelEncode):
M
MissPenguin 已提交
566
    """ Convert between text-label and text-index """
L
LDOUBLEV 已提交
567 568 569 570

    def __init__(self,
                 max_text_length,
                 character_dict_path,
文幕地方's avatar
文幕地方 已提交
571 572 573
                 replace_empty_cell_token=False,
                 merge_no_span_structure=False,
                 learn_empty_box=False,
文幕地方's avatar
fix bug  
文幕地方 已提交
574
                 point_num=2,
L
LDOUBLEV 已提交
575
                 **kwargs):
文幕地方's avatar
文幕地方 已提交
576 577 578 579 580 581 582
        self.max_text_len = max_text_length
        self.lower = False
        self.learn_empty_box = learn_empty_box
        self.merge_no_span_structure = merge_no_span_structure
        self.replace_empty_cell_token = replace_empty_cell_token

        dict_character = []
M
MissPenguin 已提交
583 584
        with open(character_dict_path, "rb") as fin:
            lines = fin.readlines()
文幕地方's avatar
文幕地方 已提交
585 586 587 588 589 590 591 592 593
            for line in lines:
                line = line.decode('utf-8').strip("\n").strip("\r\n")
                dict_character.append(line)

        dict_character = self.add_special_char(dict_character)
        self.dict = {}
        for i, char in enumerate(dict_character):
            self.dict[char] = i
        self.idx2char = {v: k for k, v in self.dict.items()}
L
LDOUBLEV 已提交
594

文幕地方's avatar
文幕地方 已提交
595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619
        self.character = dict_character
        self.point_num = point_num
        self.pad_idx = self.dict[self.beg_str]
        self.start_idx = self.dict[self.beg_str]
        self.end_idx = self.dict[self.end_str]

        self.td_token = ['<td>', '<td', '<eb></eb>', '<td></td>']
        self.empty_bbox_token_dict = {
            "[]": '<eb></eb>',
            "[' ']": '<eb1></eb1>',
            "['<b>', ' ', '</b>']": '<eb2></eb2>',
            "['\\u2028', '\\u2028']": '<eb3></eb3>',
            "['<sup>', ' ', '</sup>']": '<eb4></eb4>',
            "['<b>', '</b>']": '<eb5></eb5>',
            "['<i>', ' ', '</i>']": '<eb6></eb6>',
            "['<b>', '<i>', '</i>', '</b>']": '<eb7></eb7>',
            "['<b>', '<i>', ' ', '</i>', '</b>']": '<eb8></eb8>',
            "['<i>', '</i>']": '<eb9></eb9>',
            "['<b>', ' ', '\\u2028', ' ', '\\u2028', ' ', '</b>']":
            '<eb10></eb10>',
        }

    @property
    def _max_text_len(self):
        return self.max_text_len + 2
L
LDOUBLEV 已提交
620

M
MissPenguin 已提交
621 622
    def __call__(self, data):
        cells = data['cells']
文幕地方's avatar
文幕地方 已提交
623 624 625 626 627 628 629 630 631 632 633 634 635 636
        structure = data['structure']
        if self.merge_no_span_structure:
            structure = self._merge_no_span_structure(structure)
        if self.replace_empty_cell_token:
            structure = self._replace_empty_cell_token(structure, cells)
        # remove empty token and add " " to span token
        new_structure = []
        for token in structure:
            if token != '':
                if 'span' in token and token[0] != ' ':
                    token = ' ' + token
                new_structure.append(token)
        # encode structure
        structure = self.encode(new_structure)
M
MissPenguin 已提交
637 638
        if structure is None:
            return None
文幕地方's avatar
文幕地方 已提交
639 640 641 642 643

        structure = [self.start_idx] + structure + [self.end_idx
                                                    ]  # add sos abd eos
        structure = structure + [self.pad_idx] * (self._max_text_len -
                                                  len(structure))  # pad
M
MissPenguin 已提交
644 645 646
        structure = np.array(structure)
        data['structure'] = structure

文幕地方's avatar
文幕地方 已提交
647
        if len(structure) > self._max_text_len:
M
MissPenguin 已提交
648 649
            return None

文幕地方's avatar
文幕地方 已提交
650 651
        # encode box
        bboxes = np.zeros(
文幕地方's avatar
fix bug  
文幕地方 已提交
652
            (self._max_text_len, self.point_num * 2), dtype=np.float32)
文幕地方's avatar
文幕地方 已提交
653 654 655
        bbox_masks = np.zeros((self._max_text_len, 1), dtype=np.float32)

        bbox_idx = 0
文幕地方's avatar
fix bug  
文幕地方 已提交
656

文幕地方's avatar
文幕地方 已提交
657 658
        for i, token in enumerate(structure):
            if self.idx2char[token] in self.td_token:
文幕地方's avatar
fix bug  
文幕地方 已提交
659 660
                if 'bbox' in cells[bbox_idx] and len(cells[bbox_idx][
                        'tokens']) > 0:
文幕地方's avatar
文幕地方 已提交
661 662 663 664 665 666 667 668 669 670
                    bbox = cells[bbox_idx]['bbox'].copy()
                    bbox = np.array(bbox, dtype=np.float32).reshape(-1)
                    bboxes[i] = bbox
                    bbox_masks[i] = 1.0
                if self.learn_empty_box:
                    bbox_masks[i] = 1.0
                bbox_idx += 1
        data['bboxes'] = bboxes
        data['bbox_masks'] = bbox_masks
        return data
M
MissPenguin 已提交
671

文幕地方's avatar
文幕地方 已提交
672
    def _merge_no_span_structure(self, structure):
文幕地方's avatar
add ref  
文幕地方 已提交
673
        """
文幕地方's avatar
fix bug  
文幕地方 已提交
674
        This code is refer from:
文幕地方's avatar
add ref  
文幕地方 已提交
675 676
        https://github.com/JiaquanYe/TableMASTER-mmocr/blob/master/table_recognition/data_preprocess.py
        """
文幕地方's avatar
文幕地方 已提交
677 678 679 680 681 682 683 684 685 686 687 688
        new_structure = []
        i = 0
        while i < len(structure):
            token = structure[i]
            if token == '<td>':
                token = '<td></td>'
                i += 1
            new_structure.append(token)
            i += 1
        return new_structure

    def _replace_empty_cell_token(self, token_list, cells):
文幕地方's avatar
add ref  
文幕地方 已提交
689 690 691 692 693
        """
        This fun code is refer from:
        https://github.com/JiaquanYe/TableMASTER-mmocr/blob/master/table_recognition/data_preprocess.py
        """

文幕地方's avatar
文幕地方 已提交
694 695 696 697 698 699 700 701 702
        bbox_idx = 0
        add_empty_bbox_token_list = []
        for token in token_list:
            if token in ['<td></td>', '<td', '<td>']:
                if 'bbox' not in cells[bbox_idx].keys():
                    content = str(cells[bbox_idx]['tokens'])
                    token = self.empty_bbox_token_dict[content]
                add_empty_bbox_token_list.append(token)
                bbox_idx += 1
M
MissPenguin 已提交
703
            else:
文幕地方's avatar
文幕地方 已提交
704 705 706 707 708 709 710 711 712 713 714 715 716
                add_empty_bbox_token_list.append(token)
        return add_empty_bbox_token_list


class TableMasterLabelEncode(TableLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path,
                 replace_empty_cell_token=False,
                 merge_no_span_structure=False,
                 learn_empty_box=False,
文幕地方's avatar
fix bug  
文幕地方 已提交
717
                 point_num=2,
文幕地方's avatar
文幕地方 已提交
718 719 720 721
                 **kwargs):
        super(TableMasterLabelEncode, self).__init__(
            max_text_length, character_dict_path, replace_empty_cell_token,
            merge_no_span_structure, learn_empty_box, point_num, **kwargs)
文幕地方's avatar
fix bug  
文幕地方 已提交
722 723
        self.pad_idx = self.dict[self.pad_str]
        self.unknown_idx = self.dict[self.unknown_str]
文幕地方's avatar
文幕地方 已提交
724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768

    @property
    def _max_text_len(self):
        return self.max_text_len

    def add_special_char(self, dict_character):
        self.beg_str = '<SOS>'
        self.end_str = '<EOS>'
        self.unknown_str = '<UKN>'
        self.pad_str = '<PAD>'
        dict_character = dict_character
        dict_character = dict_character + [
            self.unknown_str, self.beg_str, self.end_str, self.pad_str
        ]
        return dict_character


class TableBoxEncode(object):
    def __init__(self, use_xywh=False, **kwargs):
        self.use_xywh = use_xywh

    def __call__(self, data):
        img_height, img_width = data['image'].shape[:2]
        bboxes = data['bboxes']
        if self.use_xywh and bboxes.shape[1] == 4:
            bboxes = self.xyxy2xywh(bboxes)
        bboxes[:, 0::2] /= img_width
        bboxes[:, 1::2] /= img_height
        data['bboxes'] = bboxes
        return data

    def xyxy2xywh(self, bboxes):
        """
        Convert coord (x1,y1,x2,y2) to (x,y,w,h).
        where (x1,y1) is top-left, (x2,y2) is bottom-right.
        (x,y) is bbox center and (w,h) is width and height.
        :param bboxes: (x1, y1, x2, y2)
        :return:
        """
        new_bboxes = np.empty_like(bboxes)
        new_bboxes[:, 0] = (bboxes[:, 0] + bboxes[:, 2]) / 2  # x center
        new_bboxes[:, 1] = (bboxes[:, 1] + bboxes[:, 3]) / 2  # y center
        new_bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]  # width
        new_bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]  # height
        return new_bboxes
A
andyjpaddle 已提交
769 770 771 772 773 774 775 776 777 778


class SARLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
779 780
        super(SARLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
A
andyjpaddle 已提交
781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805

    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 __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
        if len(text) >= self.max_text_len - 1:
            return None
        data['length'] = np.array(len(text))
        target = [self.start_idx] + text + [self.end_idx]
        padded_text = [self.padding_idx for _ in range(self.max_text_len)]
T
tink2123 已提交
806

A
andyjpaddle 已提交
807 808 809 810 811 812
        padded_text[:len(target)] = target
        data['label'] = np.array(padded_text)
        return data

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


815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861
class PRENLabelEncode(BaseRecLabelEncode):
    def __init__(self,
                 max_text_length,
                 character_dict_path,
                 use_space_char=False,
                 **kwargs):
        super(PRENLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)

    def add_special_char(self, dict_character):
        padding_str = '<PAD>'  # 0 
        end_str = '<EOS>'  # 1
        unknown_str = '<UNK>'  # 2

        dict_character = [padding_str, end_str, unknown_str] + dict_character
        self.padding_idx = 0
        self.end_idx = 1
        self.unknown_idx = 2

        return dict_character

    def encode(self, text):
        if len(text) == 0 or len(text) >= self.max_text_len:
            return None
        if self.lower:
            text = text.lower()
        text_list = []
        for char in text:
            if char not in self.dict:
                text_list.append(self.unknown_idx)
            else:
                text_list.append(self.dict[char])
        text_list.append(self.end_idx)
        if len(text_list) < self.max_text_len:
            text_list += [self.padding_idx] * (
                self.max_text_len - len(text_list))
        return text_list

    def __call__(self, data):
        text = data['label']
        encoded_text = self.encode(text)
        if encoded_text is None:
            return None
        data['label'] = np.array(encoded_text)
        return data


862 863
class VQATokenLabelEncode(object):
    """
文幕地方's avatar
文幕地方 已提交
864
    Label encode for NLP VQA methods
865 866 867 868 869 870 871
    """

    def __init__(self,
                 class_path,
                 contains_re=False,
                 add_special_ids=False,
                 algorithm='LayoutXLM',
872
                 use_textline_bbox_info=True,
873 874 875 876
                 infer_mode=False,
                 ocr_engine=None,
                 **kwargs):
        super(VQATokenLabelEncode, self).__init__()
文幕地方's avatar
文幕地方 已提交
877
        from paddlenlp.transformers import LayoutXLMTokenizer, LayoutLMTokenizer, LayoutLMv2Tokenizer
878 879 880 881 882 883 884 885 886
        from ppocr.utils.utility import load_vqa_bio_label_maps
        tokenizer_dict = {
            'LayoutXLM': {
                'class': LayoutXLMTokenizer,
                'pretrained_model': 'layoutxlm-base-uncased'
            },
            'LayoutLM': {
                'class': LayoutLMTokenizer,
                'pretrained_model': 'layoutlm-base-uncased'
文幕地方's avatar
文幕地方 已提交
887 888 889 890
            },
            'LayoutLMv2': {
                'class': LayoutLMv2Tokenizer,
                'pretrained_model': 'layoutlmv2-base-uncased'
891 892 893 894 895 896 897 898 899 900
            }
        }
        self.contains_re = contains_re
        tokenizer_config = tokenizer_dict[algorithm]
        self.tokenizer = tokenizer_config['class'].from_pretrained(
            tokenizer_config['pretrained_model'])
        self.label2id_map, id2label_map = load_vqa_bio_label_maps(class_path)
        self.add_special_ids = add_special_ids
        self.infer_mode = infer_mode
        self.ocr_engine = ocr_engine
901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935
        self.use_textline_bbox_info = use_textline_bbox_info

    def split_bbox(self, bbox, text, tokenizer):
        words = text.split()
        token_bboxes = []
        curr_word_idx = 0
        x1, y1, x2, y2 = bbox
        unit_w = (x2 - x1) / len(text)
        for idx, word in enumerate(words):
            curr_w = len(word) * unit_w
            word_bbox = [x1, y1, x1 + curr_w, y2]
            token_bboxes.extend([word_bbox] * len(tokenizer.tokenize(word)))
            x1 += (len(word) + 1) * unit_w
        return token_bboxes

    def filter_empty_contents(self, ocr_info):
        """
        find out the empty texts and remove the links
        """
        new_ocr_info = []
        empty_index = []
        for idx, info in enumerate(ocr_info):
            if len(info["transcription"]) > 0:
                new_ocr_info.append(copy.deepcopy(info))
            else:
                empty_index.append(info["id"])

        for idx, info in enumerate(new_ocr_info):
            new_link = []
            for link in info["linking"]:
                if link[0] in empty_index or link[1] in empty_index:
                    continue
                new_link.append(link)
            new_ocr_info[idx]["linking"] = new_link
        return new_ocr_info
936 937

    def __call__(self, data):
文幕地方's avatar
文幕地方 已提交
938 939
        # load bbox and label info
        ocr_info = self._load_ocr_info(data)
940

941 942 943 944 945
        # for re
        train_re = self.contains_re and not self.infer_mode
        if train_re:
            ocr_info = self.filter_empty_contents(ocr_info)

文幕地方's avatar
文幕地方 已提交
946
        height, width, _ = data['image'].shape
947 948 949 950 951

        words_list = []
        bbox_list = []
        input_ids_list = []
        token_type_ids_list = []
文幕地方's avatar
文幕地方 已提交
952
        segment_offset_id = []
953 954
        gt_label_list = []

文幕地方's avatar
文幕地方 已提交
955 956 957 958 959 960 961
        entities = []

        if train_re:
            relations = []
            id2label = {}
            entity_id_to_index_map = {}
            empty_entity = set()
文幕地方's avatar
文幕地方 已提交
962 963 964 965

        data['ocr_info'] = copy.deepcopy(ocr_info)

        for info in ocr_info:
966 967 968
            text = info["transcription"]
            if len(text) <= 0:
                continue
文幕地方's avatar
文幕地方 已提交
969
            if train_re:
970
                # for re
971
                if len(text) == 0:
972 973 974 975
                    empty_entity.add(info["id"])
                    continue
                id2label[info["id"]] = info["label"]
                relations.extend([tuple(sorted(l)) for l in info["linking"]])
文幕地方's avatar
文幕地方 已提交
976
            # smooth_box
977
            info["bbox"] = self.trans_poly_to_bbox(info["points"])
978 979 980 981 982 983 984 985 986 987 988

            encode_res = self.tokenizer.encode(
                text, pad_to_max_seq_len=False, return_attention_mask=True)

            if not self.add_special_ids:
                # TODO: use tok.all_special_ids to remove
                encode_res["input_ids"] = encode_res["input_ids"][1:-1]
                encode_res["token_type_ids"] = encode_res["token_type_ids"][1:
                                                                            -1]
                encode_res["attention_mask"] = encode_res["attention_mask"][1:
                                                                            -1]
989 990 991 992 993 994 995 996 997 998 999 1000 1001

            if self.use_textline_bbox_info:
                bbox = [info["bbox"]] * len(encode_res["input_ids"])
            else:
                bbox = self.split_bbox(info["bbox"], info["transcription"],
                                       self.tokenizer)
            if len(bbox) <= 0:
                continue
            bbox = self._smooth_box(bbox, height, width)
            if self.add_special_ids:
                bbox.insert(0, [0, 0, 0, 0])
                bbox.append([0, 0, 0, 0])

文幕地方's avatar
文幕地方 已提交
1002 1003 1004 1005 1006 1007
            # parse label
            if not self.infer_mode:
                label = info['label']
                gt_label = self._parse_label(label, encode_res)

            # construct entities for re
文幕地方's avatar
文幕地方 已提交
1008 1009 1010 1011
            if train_re:
                if gt_label[0] != self.label2id_map["O"]:
                    entity_id_to_index_map[info["id"]] = len(entities)
                    label = label.upper()
1012 1013 1014 1015
                    entities.append({
                        "start": len(input_ids_list),
                        "end":
                        len(input_ids_list) + len(encode_res["input_ids"]),
文幕地方's avatar
文幕地方 已提交
1016
                        "label": label.upper(),
1017
                    })
文幕地方's avatar
文幕地方 已提交
1018 1019 1020 1021 1022 1023
            else:
                entities.append({
                    "start": len(input_ids_list),
                    "end": len(input_ids_list) + len(encode_res["input_ids"]),
                    "label": 'O',
                })
1024 1025
            input_ids_list.extend(encode_res["input_ids"])
            token_type_ids_list.extend(encode_res["token_type_ids"])
1026
            bbox_list.extend(bbox)
1027
            words_list.append(text)
文幕地方's avatar
文幕地方 已提交
1028 1029 1030 1031 1032 1033 1034 1035 1036 1037
            segment_offset_id.append(len(input_ids_list))
            if not self.infer_mode:
                gt_label_list.extend(gt_label)

        data['input_ids'] = input_ids_list
        data['token_type_ids'] = token_type_ids_list
        data['bbox'] = bbox_list
        data['attention_mask'] = [1] * len(input_ids_list)
        data['labels'] = gt_label_list
        data['segment_offset_id'] = segment_offset_id
1038 1039 1040 1041
        data['tokenizer_params'] = dict(
            padding_side=self.tokenizer.padding_side,
            pad_token_type_id=self.tokenizer.pad_token_type_id,
            pad_token_id=self.tokenizer.pad_token_id)
文幕地方's avatar
文幕地方 已提交
1042
        data['entities'] = entities
1043

文幕地方's avatar
文幕地方 已提交
1044 1045 1046 1047 1048
        if train_re:
            data['relations'] = relations
            data['id2label'] = id2label
            data['empty_entity'] = empty_entity
            data['entity_id_to_index_map'] = entity_id_to_index_map
1049 1050
        return data

1051 1052 1053 1054 1055 1056
    def trans_poly_to_bbox(self, poly):
        x1 = np.min([p[0] for p in poly])
        x2 = np.max([p[0] for p in poly])
        y1 = np.min([p[1] for p in poly])
        y2 = np.max([p[1] for p in poly])
        return [x1, y1, x2, y2]
文幕地方's avatar
文幕地方 已提交
1057

1058
    def _load_ocr_info(self, data):
文幕地方's avatar
文幕地方 已提交
1059 1060 1061 1062 1063
        if self.infer_mode:
            ocr_result = self.ocr_engine.ocr(data['image'], cls=False)
            ocr_info = []
            for res in ocr_result:
                ocr_info.append({
1064 1065 1066
                    "transcription": res[1][0],
                    "bbox": self.trans_poly_to_bbox(res[0]),
                    "points": res[0],
文幕地方's avatar
文幕地方 已提交
1067 1068 1069 1070 1071 1072
                })
            return ocr_info
        else:
            info = data['label']
            # read text info
            info_dict = json.loads(info)
1073
            return info_dict
文幕地方's avatar
文幕地方 已提交
1074

1075 1076 1077 1078 1079 1080 1081 1082
    def _smooth_box(self, bboxes, height, width):
        bboxes = np.array(bboxes)
        bboxes[:, 0] = bboxes[:, 0] * 1000 / width
        bboxes[:, 2] = bboxes[:, 2] * 1000 / width
        bboxes[:, 1] = bboxes[:, 1] * 1000 / height
        bboxes[:, 3] = bboxes[:, 3] * 1000 / height
        bboxes = bboxes.astype("int64").tolist()
        return bboxes
文幕地方's avatar
文幕地方 已提交
1083 1084 1085

    def _parse_label(self, label, encode_res):
        gt_label = []
1086
        if label.lower() in ["other", "others", "ignore"]:
文幕地方's avatar
文幕地方 已提交
1087 1088 1089 1090 1091 1092
            gt_label.extend([0] * len(encode_res["input_ids"]))
        else:
            gt_label.append(self.label2id_map[("b-" + label).upper()])
            gt_label.extend([self.label2id_map[("i-" + label).upper()]] *
                            (len(encode_res["input_ids"]) - 1))
        return gt_label
A
andyjpaddle 已提交
1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122


class MultiLabelEncode(BaseRecLabelEncode):
    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
        super(MultiLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)

        self.ctc_encode = CTCLabelEncode(max_text_length, character_dict_path,
                                         use_space_char, **kwargs)
        self.sar_encode = SARLabelEncode(max_text_length, character_dict_path,
                                         use_space_char, **kwargs)

    def __call__(self, data):
        data_ctc = copy.deepcopy(data)
        data_sar = copy.deepcopy(data)
        data_out = dict()
        data_out['img_path'] = data.get('img_path', None)
        data_out['image'] = data['image']
        ctc = self.ctc_encode.__call__(data_ctc)
        sar = self.sar_encode.__call__(data_sar)
        if ctc is None or sar is None:
            return None
        data_out['label_ctc'] = ctc['label']
        data_out['label_sar'] = sar['label']
        data_out['length'] = ctc['length']
        return data_out
1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218


class NRTRLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):

        super(NRTRLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)

    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
        if len(text) >= self.max_text_len - 1:
            return None
        data['length'] = np.array(len(text))
        text.insert(0, 2)
        text.append(3)
        text = text + [0] * (self.max_text_len - len(text))
        data['label'] = np.array(text)
        return data

    def add_special_char(self, dict_character):
        dict_character = ['blank', '<unk>', '<s>', '</s>'] + dict_character
        return dict_character


class ViTSTRLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 ignore_index=0,
                 **kwargs):

        super(ViTSTRLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
        self.ignore_index = ignore_index

    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
        if len(text) >= self.max_text_len:
            return None
        data['length'] = np.array(len(text))
        text.insert(0, self.ignore_index)
        text.append(1)
        text = text + [self.ignore_index] * (self.max_text_len + 2 - len(text))
        data['label'] = np.array(text)
        return data

    def add_special_char(self, dict_character):
        dict_character = ['<s>', '</s>'] + dict_character
        return dict_character


class ABINetLabelEncode(BaseRecLabelEncode):
    """ Convert between text-label and text-index """

    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 ignore_index=100,
                 **kwargs):

        super(ABINetLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
        self.ignore_index = ignore_index

    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
        if len(text) >= self.max_text_len:
            return None
        data['length'] = np.array(len(text))
        text.append(0)
        text = text + [self.ignore_index] * (self.max_text_len + 1 - len(text))
        data['label'] = np.array(text)
        return data

    def add_special_char(self, dict_character):
        dict_character = ['</s>'] + dict_character
        return dict_character