label_ops.py 27.1 KB
Newer Older
WenmuZhou's avatar
WenmuZhou committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 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

import numpy as np
tink2123's avatar
tink2123 committed
21
import string
LDOUBLEV's avatar
add kie  
LDOUBLEV committed
22
from shapely.geometry import LineString, Point, Polygon
LDOUBLEV's avatar
LDOUBLEV committed
23
import json
WenmuZhou's avatar
WenmuZhou committed
24

tink2123's avatar
tink2123 committed
25
26
from ppocr.utils.logging import get_logger

WenmuZhou's avatar
WenmuZhou committed
27
28
29
30
31
32
33
34
35
36
37
38

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
WenmuZhou's avatar
WenmuZhou committed
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58


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)
LDOUBLEV's avatar
LDOUBLEV committed
59
60
        if len(boxes) == 0:
            return None
MissPenguin's avatar
MissPenguin committed
61
        boxes = self.expand_points_num(boxes)
WenmuZhou's avatar
WenmuZhou committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
        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)]
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]
        rect[3] = pts[np.argmax(diff)]
        return rect

MissPenguin's avatar
MissPenguin committed
80
81
82
83
84
85
86
87
88
89
90
    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

WenmuZhou's avatar
WenmuZhou committed
91
92
93
94
95
96
97
98
99
100

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
tink2123's avatar
tink2123 committed
101
102
        self.beg_str = "sos"
        self.end_str = "eos"
tink2123's avatar
tink2123 committed
103
        self.lower = False
tink2123's avatar
tink2123 committed
104
105
106
107
108
109

        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"
            )
WenmuZhou's avatar
WenmuZhou committed
110
111
            self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz"
            dict_character = list(self.character_str)
tink2123's avatar
tink2123 committed
112
113
            self.lower = True
        else:
WenmuZhou's avatar
WenmuZhou committed
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
            self.character_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 += line
            if use_space_char:
                self.character_str += " "
            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]
        """
WenmuZhou's avatar
WenmuZhou committed
142
        if len(text) == 0 or len(text) > self.max_text_len:
WenmuZhou's avatar
WenmuZhou committed
143
            return None
tink2123's avatar
tink2123 committed
144
        if self.lower:
WenmuZhou's avatar
WenmuZhou committed
145
146
147
148
149
150
151
152
153
154
155
156
157
            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


Topdu's avatar
Topdu committed
158
159
160
161
162
163
164
165
166
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):

tink2123's avatar
tink2123 committed
167
168
        super(NRTRLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
tink2123's avatar
tink2123 committed
169

Topdu's avatar
Topdu committed
170
171
172
173
174
    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
Topdu's avatar
Topdu committed
175
176
        if len(text) >= self.max_text_len - 1:
            return None
Topdu's avatar
Topdu committed
177
178
179
180
181
182
        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
tink2123's avatar
tink2123 committed
183

Topdu's avatar
Topdu committed
184
    def add_special_char(self, dict_character):
tink2123's avatar
tink2123 committed
185
        dict_character = ['blank', '<unk>', '<s>', '</s>'] + dict_character
Topdu's avatar
Topdu committed
186
187
        return dict_character

tink2123's avatar
tink2123 committed
188

WenmuZhou's avatar
WenmuZhou committed
189
190
191
192
193
194
195
196
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):
tink2123's avatar
tink2123 committed
197
198
        super(CTCLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
WenmuZhou's avatar
WenmuZhou committed
199
200
201
202
203
204
205
206
207

    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)
208
209
210
211
212

        label = [0] * len(self.character)
        for x in text:
            label[x] += 1
        data['label_ace'] = np.array(label)
WenmuZhou's avatar
WenmuZhou committed
213
214
215
216
217
218
219
        return data

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


Jethong's avatar
Jethong committed
220
class E2ELabelEncodeTest(BaseRecLabelEncode):
Jethong's avatar
Jethong committed
221
222
223
224
225
    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
tink2123's avatar
tink2123 committed
226
227
        super(E2ELabelEncodeTest, self).__init__(
            max_text_length, character_dict_path, use_space_char)
Jethong's avatar
Jethong committed
228
229

    def __call__(self, data):
Jethong's avatar
Jethong committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
        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
Jethong's avatar
Jethong committed
248
        data['ignore_tags'] = txt_tags
Jethong's avatar
Jethong committed
249
        temp_texts = []
Jethong's avatar
Jethong committed
250
        for text in txts:
Jethong's avatar
Jethong committed
251
            text = text.lower()
Jethong's avatar
Jethong committed
252
253
254
            text = self.encode(text)
            if text is None:
                return None
Jethong's avatar
Jethong committed
255
256
            text = text + [padnum] * (self.max_text_len - len(text)
                                      )  # use 36 to pad
Jethong's avatar
Jethong committed
257
258
259
260
261
            temp_texts.append(text)
        data['texts'] = np.array(temp_texts)
        return data


Jethong's avatar
Jethong committed
262
class E2ELabelEncodeTrain(object):
Jethong's avatar
Jethong committed
263
264
    def __init__(self, **kwargs):
        pass
Jethong's avatar
Jethong committed
265
266

    def __call__(self, data):
Jethong's avatar
Jethong committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
        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
Jethong's avatar
Jethong committed
286
        data['ignore_tags'] = txt_tags
Jethong's avatar
Jethong committed
287
288
289
        return data


LDOUBLEV's avatar
add kie  
LDOUBLEV committed
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
class KieLabelEncode(object):
    def __init__(self, character_dict_path, norm=10, directed=False, **kwargs):
        super(KieLabelEncode, self).__init__()
        self.dict = dict({'': 0})
        with open(character_dict_path, 'r') as fr:
            idx = 1
            for line in fr:
                char = line.strip()
                self.dict[char] = idx
                idx += 1
        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."""
LDOUBLEV's avatar
debug  
LDOUBLEV committed
318
        max_len = 300
LDOUBLEV's avatar
add kie  
LDOUBLEV committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
        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)
LDOUBLEV's avatar
debug  
LDOUBLEV committed
344
        max_num = 300
LDOUBLEV's avatar
add kie  
LDOUBLEV committed
345
346
347
348
349
350
351
        temp_bboxes = np.zeros([max_num, 4])
        h, _ = bboxes.shape
        temp_bboxes[:h, :h] = bboxes

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

LDOUBLEV's avatar
debug  
LDOUBLEV committed
352
        temp_padded_text_inds = np.zeros([max_num, max_num])
LDOUBLEV's avatar
add kie  
LDOUBLEV committed
353
354
        temp_padded_text_inds[:h, :] = padded_text_inds

LDOUBLEV's avatar
debug  
LDOUBLEV committed
355
        temp_labels = np.zeros([max_num, max_num])
LDOUBLEV's avatar
add kie  
LDOUBLEV committed
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
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
        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)
            labels.append(ann['label'])
            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)


WenmuZhou's avatar
WenmuZhou committed
452
453
454
455
456
457
458
459
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):
tink2123's avatar
tink2123 committed
460
461
        super(AttnLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
WenmuZhou's avatar
WenmuZhou committed
462
463

    def add_special_char(self, dict_character):
LDOUBLEV's avatar
LDOUBLEV committed
464
465
466
        self.beg_str = "sos"
        self.end_str = "eos"
        dict_character = [self.beg_str] + dict_character + [self.end_str]
WenmuZhou's avatar
WenmuZhou committed
467
468
        return dict_character

LDOUBLEV's avatar
LDOUBLEV committed
469
470
    def __call__(self, data):
        text = data['label']
WenmuZhou's avatar
WenmuZhou committed
471
        text = self.encode(text)
LDOUBLEV's avatar
LDOUBLEV committed
472
473
        if text is None:
            return None
LDOUBLEV's avatar
LDOUBLEV committed
474
        if len(text) >= self.max_text_len:
LDOUBLEV's avatar
LDOUBLEV committed
475
476
477
            return None
        data['length'] = np.array(len(text))
        text = [0] + text + [len(self.character) - 1] + [0] * (self.max_text_len
tink2123's avatar
tink2123 committed
478
                                                               - len(text) - 2)
LDOUBLEV's avatar
LDOUBLEV committed
479
480
481
482
483
484
485
        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]
WenmuZhou's avatar
WenmuZhou committed
486
487
488
489
490
491
492
493
494
495

    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
tink2123's avatar
tink2123 committed
496
497


tink2123's avatar
tink2123 committed
498
499
500
501
502
503
504
505
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):
tink2123's avatar
tink2123 committed
506
507
        super(SEEDLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
tink2123's avatar
tink2123 committed
508
509
510
511
512
513
514
515
516
517
518
519
520

    def add_special_char(self, dict_character):
        self.end_str = "eos"
        dict_character = dict_character + [self.end_str]
        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
tink2123's avatar
rm anno  
tink2123 committed
521
        data['length'] = np.array(len(text)) + 1  # conclude eos
tink2123's avatar
tink2123 committed
522
523
524
525
526
527
        text = text + [len(self.character) - 1] * (self.max_text_len - len(text)
                                                   )
        data['label'] = np.array(text)
        return data


tink2123's avatar
tink2123 committed
528
529
530
531
532
533
534
535
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):
tink2123's avatar
tink2123 committed
536
537
        super(SRNLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
tink2123's avatar
tink2123 committed
538
539
540
541
542
543
544
545

    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)
tink2123's avatar
tink2123 committed
546
        char_num = len(self.character)
tink2123's avatar
tink2123 committed
547
548
549
550
551
        if text is None:
            return None
        if len(text) > self.max_text_len:
            return None
        data['length'] = np.array(len(text))
tink2123's avatar
tink2123 committed
552
        text = text + [char_num - 1] * (self.max_text_len - len(text))
tink2123's avatar
tink2123 committed
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
        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
MissPenguin's avatar
MissPenguin committed
570

LDOUBLEV's avatar
LDOUBLEV committed
571

MissPenguin's avatar
MissPenguin committed
572
573
class TableLabelEncode(object):
    """ Convert between text-label and text-index """
LDOUBLEV's avatar
LDOUBLEV committed
574
575
576
577
578
579
580
581

    def __init__(self,
                 max_text_length,
                 max_elem_length,
                 max_cell_num,
                 character_dict_path,
                 span_weight=1.0,
                 **kwargs):
MissPenguin's avatar
MissPenguin committed
582
583
584
        self.max_text_length = max_text_length
        self.max_elem_length = max_elem_length
        self.max_cell_num = max_cell_num
LDOUBLEV's avatar
LDOUBLEV committed
585
586
        list_character, list_elem = self.load_char_elem_dict(
            character_dict_path)
MissPenguin's avatar
MissPenguin committed
587
588
589
590
591
592
593
594
595
        list_character = self.add_special_char(list_character)
        list_elem = self.add_special_char(list_elem)
        self.dict_character = {}
        for i, char in enumerate(list_character):
            self.dict_character[char] = i
        self.dict_elem = {}
        for i, elem in enumerate(list_elem):
            self.dict_elem[elem] = i
        self.span_weight = span_weight
LDOUBLEV's avatar
LDOUBLEV committed
596

MissPenguin's avatar
MissPenguin committed
597
598
599
600
601
    def load_char_elem_dict(self, character_dict_path):
        list_character = []
        list_elem = []
        with open(character_dict_path, "rb") as fin:
            lines = fin.readlines()
WenmuZhou's avatar
WenmuZhou committed
602
            substr = lines[0].decode('utf-8').strip("\r\n").split("\t")
MissPenguin's avatar
MissPenguin committed
603
604
            character_num = int(substr[0])
            elem_num = int(substr[1])
LDOUBLEV's avatar
LDOUBLEV committed
605
            for cno in range(1, 1 + character_num):
WenmuZhou's avatar
WenmuZhou committed
606
                character = lines[cno].decode('utf-8').strip("\r\n")
MissPenguin's avatar
MissPenguin committed
607
                list_character.append(character)
LDOUBLEV's avatar
LDOUBLEV committed
608
            for eno in range(1 + character_num, 1 + character_num + elem_num):
WenmuZhou's avatar
WenmuZhou committed
609
                elem = lines[eno].decode('utf-8').strip("\r\n")
MissPenguin's avatar
MissPenguin committed
610
611
                list_elem.append(elem)
        return list_character, list_elem
LDOUBLEV's avatar
LDOUBLEV committed
612

MissPenguin's avatar
MissPenguin committed
613
614
615
616
617
    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
LDOUBLEV's avatar
LDOUBLEV committed
618

MissPenguin's avatar
MissPenguin committed
619
620
621
622
623
624
    def get_span_idx_list(self):
        span_idx_list = []
        for elem in self.dict_elem:
            if 'span' in elem:
                span_idx_list.append(self.dict_elem[elem])
        return span_idx_list
LDOUBLEV's avatar
LDOUBLEV committed
625

MissPenguin's avatar
MissPenguin committed
626
627
628
629
630
631
632
633
    def __call__(self, data):
        cells = data['cells']
        structure = data['structure']['tokens']
        structure = self.encode(structure, 'elem')
        if structure is None:
            return None
        elem_num = len(structure)
        structure = [0] + structure + [len(self.dict_elem) - 1]
LDOUBLEV's avatar
LDOUBLEV committed
634
635
        structure = structure + [0] * (self.max_elem_length + 2 - len(structure)
                                       )
MissPenguin's avatar
MissPenguin committed
636
637
638
639
640
        structure = np.array(structure)
        data['structure'] = structure
        elem_char_idx1 = self.dict_elem['<td>']
        elem_char_idx2 = self.dict_elem['<td']
        span_idx_list = self.get_span_idx_list()
LDOUBLEV's avatar
LDOUBLEV committed
641
642
        td_idx_list = np.logical_or(structure == elem_char_idx1,
                                    structure == elem_char_idx2)
MissPenguin's avatar
MissPenguin committed
643
        td_idx_list = np.where(td_idx_list)[0]
LDOUBLEV's avatar
LDOUBLEV committed
644
645
646

        structure_mask = np.ones(
            (self.max_elem_length + 2, 1), dtype=np.float32)
MissPenguin's avatar
MissPenguin committed
647
        bbox_list = np.zeros((self.max_elem_length + 2, 4), dtype=np.float32)
LDOUBLEV's avatar
LDOUBLEV committed
648
649
        bbox_list_mask = np.zeros(
            (self.max_elem_length + 2, 1), dtype=np.float32)
MissPenguin's avatar
MissPenguin committed
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
        img_height, img_width, img_ch = data['image'].shape
        if len(span_idx_list) > 0:
            span_weight = len(td_idx_list) * 1.0 / len(span_idx_list)
            span_weight = min(max(span_weight, 1.0), self.span_weight)
        for cno in range(len(cells)):
            if 'bbox' in cells[cno]:
                bbox = cells[cno]['bbox'].copy()
                bbox[0] = bbox[0] * 1.0 / img_width
                bbox[1] = bbox[1] * 1.0 / img_height
                bbox[2] = bbox[2] * 1.0 / img_width
                bbox[3] = bbox[3] * 1.0 / img_height
                td_idx = td_idx_list[cno]
                bbox_list[td_idx] = bbox
                bbox_list_mask[td_idx] = 1.0
                cand_span_idx = td_idx + 1
                if cand_span_idx < (self.max_elem_length + 2):
                    if structure[cand_span_idx] in span_idx_list:
                        structure_mask[cand_span_idx] = span_weight

        data['bbox_list'] = bbox_list
        data['bbox_list_mask'] = bbox_list_mask
        data['structure_mask'] = structure_mask
        char_beg_idx = self.get_beg_end_flag_idx('beg', 'char')
        char_end_idx = self.get_beg_end_flag_idx('end', 'char')
        elem_beg_idx = self.get_beg_end_flag_idx('beg', 'elem')
        elem_end_idx = self.get_beg_end_flag_idx('end', 'elem')
LDOUBLEV's avatar
LDOUBLEV committed
676
677
678
679
680
        data['sp_tokens'] = np.array([
            char_beg_idx, char_end_idx, elem_beg_idx, elem_end_idx,
            elem_char_idx1, elem_char_idx2, self.max_text_length,
            self.max_elem_length, self.max_cell_num, elem_num
        ])
MissPenguin's avatar
MissPenguin committed
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
        return data

    def encode(self, text, char_or_elem):
        """convert text-label into text-index.
        """
        if char_or_elem == "char":
            max_len = self.max_text_length
            current_dict = self.dict_character
        else:
            max_len = self.max_elem_length
            current_dict = self.dict_elem
        if len(text) > max_len:
            return None
        if len(text) == 0:
            if char_or_elem == "char":
                return [self.dict_character['space']]
            else:
                return None
        text_list = []
        for char in text:
            if char not in current_dict:
                return None
            text_list.append(current_dict[char])
        if len(text_list) == 0:
            if char_or_elem == "char":
                return [self.dict_character['space']]
            else:
                return None
        return text_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 = np.array(self.dict_character[self.beg_str])
            elif beg_or_end == "end":
                idx = np.array(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 = np.array(self.dict_elem[self.beg_str])
            elif beg_or_end == "end":
                idx = np.array(self.dict_elem[self.end_str])
            else:
                assert False, "Unsupport type %s in get_beg_end_flag_idx of elem" \
LDOUBLEV's avatar
LDOUBLEV committed
732
                              % beg_or_end
MissPenguin's avatar
MissPenguin committed
733
734
        else:
            assert False, "Unsupport type %s in char_or_elem" \
LDOUBLEV's avatar
LDOUBLEV committed
735
                              % char_or_elem
MissPenguin's avatar
MissPenguin committed
736
        return idx
andyjpaddle's avatar
andyjpaddle committed
737
738
739
740
741
742
743
744
745
746


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):
tink2123's avatar
tink2123 committed
747
748
        super(SARLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
andyjpaddle's avatar
andyjpaddle committed
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773

    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)]
tink2123's avatar
tink2123 committed
774

andyjpaddle's avatar
andyjpaddle committed
775
776
777
778
779
780
        padded_text[:len(target)] = target
        data['label'] = np.array(padded_text)
        return data

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