label_ops.py 27.2 KB
Newer Older
W
WenmuZhou 已提交
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
T
tink2123 已提交
21
import string
L
add kie  
LDOUBLEV 已提交
22
from shapely.geometry import LineString, Point, Polygon
L
LDOUBLEV 已提交
23
import json
W
WenmuZhou 已提交
24

T
tink2123 已提交
25 26
from ppocr.utils.logging import get_logger

W
WenmuZhou 已提交
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
W
WenmuZhou 已提交
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)
L
LDOUBLEV 已提交
59 60
        if len(boxes) == 0:
            return None
M
MissPenguin 已提交
61
        boxes = self.expand_points_num(boxes)
W
WenmuZhou 已提交
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

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

W
WenmuZhou 已提交
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
T
tink2123 已提交
101 102
        self.beg_str = "sos"
        self.end_str = "eos"
T
tink2123 已提交
103
        self.lower = False
T
tink2123 已提交
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"
            )
W
WenmuZhou 已提交
110 111
            self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz"
            dict_character = list(self.character_str)
T
tink2123 已提交
112 113
            self.lower = True
        else:
W
WenmuZhou 已提交
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]
        """
W
WenmuZhou 已提交
142
        if len(text) == 0 or len(text) > self.max_text_len:
W
WenmuZhou 已提交
143
            return None
T
tink2123 已提交
144
        if self.lower:
W
WenmuZhou 已提交
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


T
Topdu 已提交
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):

T
tink2123 已提交
167 168
        super(NRTRLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
T
tink2123 已提交
169

T
Topdu 已提交
170 171 172 173 174
    def __call__(self, data):
        text = data['label']
        text = self.encode(text)
        if text is None:
            return None
T
Topdu 已提交
175 176
        if len(text) >= self.max_text_len - 1:
            return None
T
Topdu 已提交
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
T
tink2123 已提交
183

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

T
tink2123 已提交
188

W
WenmuZhou 已提交
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):
T
tink2123 已提交
197 198
        super(CTCLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
W
WenmuZhou 已提交
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)
W
WenmuZhou 已提交
213 214 215 216 217 218 219
        return data

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


J
Jethong 已提交
220
class E2ELabelEncodeTest(BaseRecLabelEncode):
J
Jethong 已提交
221 222 223 224 225
    def __init__(self,
                 max_text_length,
                 character_dict_path=None,
                 use_space_char=False,
                 **kwargs):
T
tink2123 已提交
226 227
        super(E2ELabelEncodeTest, self).__init__(
            max_text_length, character_dict_path, use_space_char)
J
Jethong 已提交
228 229

    def __call__(self, data):
J
Jethong 已提交
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
J
Jethong 已提交
248
        data['ignore_tags'] = txt_tags
J
Jethong 已提交
249
        temp_texts = []
J
Jethong 已提交
250
        for text in txts:
J
Jethong 已提交
251
            text = text.lower()
J
Jethong 已提交
252 253 254
            text = self.encode(text)
            if text is None:
                return None
J
Jethong 已提交
255 256
            text = text + [padnum] * (self.max_text_len - len(text)
                                      )  # use 36 to pad
J
Jethong 已提交
257 258 259 260 261
            temp_texts.append(text)
        data['texts'] = np.array(temp_texts)
        return data


J
Jethong 已提交
262
class E2ELabelEncodeTrain(object):
J
Jethong 已提交
263 264
    def __init__(self, **kwargs):
        pass
J
Jethong 已提交
265 266

    def __call__(self, data):
J
Jethong 已提交
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
J
Jethong 已提交
286
        data['ignore_tags'] = txt_tags
J
Jethong 已提交
287 288 289
        return data


L
add kie  
LDOUBLEV 已提交
290 291 292 293
class KieLabelEncode(object):
    def __init__(self, character_dict_path, norm=10, directed=False, **kwargs):
        super(KieLabelEncode, self).__init__()
        self.dict = dict({'': 0})
L
fix win  
LDOUBLEV 已提交
294
        with open(character_dict_path, 'r', encoding='utf-8') as fr:
L
add kie  
LDOUBLEV 已提交
295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317
            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."""
L
debug  
LDOUBLEV 已提交
318
        max_len = 300
L
add kie  
LDOUBLEV 已提交
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)
L
debug  
LDOUBLEV 已提交
344
        max_num = 300
L
add kie  
LDOUBLEV 已提交
345 346
        temp_bboxes = np.zeros([max_num, 4])
        h, _ = bboxes.shape
那珈落's avatar
那珈落 已提交
347
        temp_bboxes[:h, :] = bboxes
L
add kie  
LDOUBLEV 已提交
348 349 350 351

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

L
debug  
LDOUBLEV 已提交
352
        temp_padded_text_inds = np.zeros([max_num, max_num])
L
add kie  
LDOUBLEV 已提交
353 354
        temp_padded_text_inds[:h, :] = padded_text_inds

L
debug  
LDOUBLEV 已提交
355
        temp_labels = np.zeros([max_num, max_num])
L
add kie  
LDOUBLEV 已提交
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)


W
WenmuZhou 已提交
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):
T
tink2123 已提交
460 461
        super(AttnLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
W
WenmuZhou 已提交
462 463

    def add_special_char(self, dict_character):
L
LDOUBLEV 已提交
464 465 466
        self.beg_str = "sos"
        self.end_str = "eos"
        dict_character = [self.beg_str] + dict_character + [self.end_str]
W
WenmuZhou 已提交
467 468
        return dict_character

L
LDOUBLEV 已提交
469 470
    def __call__(self, data):
        text = data['label']
W
WenmuZhou 已提交
471
        text = self.encode(text)
L
LDOUBLEV 已提交
472 473
        if text is None:
            return None
L
LDOUBLEV 已提交
474
        if len(text) >= self.max_text_len:
L
LDOUBLEV 已提交
475 476 477
            return None
        data['length'] = np.array(len(text))
        text = [0] + text + [len(self.character) - 1] + [0] * (self.max_text_len
T
tink2123 已提交
478
                                                               - len(text) - 2)
L
LDOUBLEV 已提交
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]
W
WenmuZhou 已提交
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
T
tink2123 已提交
496 497


T
tink2123 已提交
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):
T
tink2123 已提交
506 507
        super(SEEDLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
T
tink2123 已提交
508 509

    def add_special_char(self, dict_character):
T
tink2123 已提交
510
        self.padding = "padding"
T
tink2123 已提交
511
        self.end_str = "eos"
T
tink2123 已提交
512 513 514 515
        self.unknown = "unknown"
        dict_character = dict_character + [
            self.end_str, self.padding, self.unknown
        ]
T
tink2123 已提交
516 517 518 519 520 521 522 523 524
        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 已提交
525
        data['length'] = np.array(len(text)) + 1  # conclude eos
T
tink2123 已提交
526 527
        text = text + [len(self.character) - 3] + [len(self.character) - 2] * (
            self.max_text_len - len(text) - 1)
T
tink2123 已提交
528 529 530 531
        data['label'] = np.array(text)
        return data


T
tink2123 已提交
532 533 534 535 536 537 538 539
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 已提交
540 541
        super(SRNLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
T
tink2123 已提交
542 543 544 545 546 547 548 549

    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 已提交
550
        char_num = len(self.character)
T
tink2123 已提交
551 552 553 554 555
        if text is None:
            return None
        if len(text) > self.max_text_len:
            return None
        data['length'] = np.array(len(text))
T
tink2123 已提交
556
        text = text + [char_num - 1] * (self.max_text_len - len(text))
T
tink2123 已提交
557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573
        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 已提交
574

L
LDOUBLEV 已提交
575

M
MissPenguin 已提交
576 577
class TableLabelEncode(object):
    """ Convert between text-label and text-index """
L
LDOUBLEV 已提交
578 579 580 581 582 583 584 585

    def __init__(self,
                 max_text_length,
                 max_elem_length,
                 max_cell_num,
                 character_dict_path,
                 span_weight=1.0,
                 **kwargs):
M
MissPenguin 已提交
586 587 588
        self.max_text_length = max_text_length
        self.max_elem_length = max_elem_length
        self.max_cell_num = max_cell_num
L
LDOUBLEV 已提交
589 590
        list_character, list_elem = self.load_char_elem_dict(
            character_dict_path)
M
MissPenguin 已提交
591 592 593 594 595 596 597 598 599
        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
L
LDOUBLEV 已提交
600

M
MissPenguin 已提交
601 602 603 604 605
    def load_char_elem_dict(self, character_dict_path):
        list_character = []
        list_elem = []
        with open(character_dict_path, "rb") as fin:
            lines = fin.readlines()
W
WenmuZhou 已提交
606
            substr = lines[0].decode('utf-8').strip("\r\n").split("\t")
M
MissPenguin 已提交
607 608
            character_num = int(substr[0])
            elem_num = int(substr[1])
L
LDOUBLEV 已提交
609
            for cno in range(1, 1 + character_num):
W
WenmuZhou 已提交
610
                character = lines[cno].decode('utf-8').strip("\r\n")
M
MissPenguin 已提交
611
                list_character.append(character)
L
LDOUBLEV 已提交
612
            for eno in range(1 + character_num, 1 + character_num + elem_num):
W
WenmuZhou 已提交
613
                elem = lines[eno].decode('utf-8').strip("\r\n")
M
MissPenguin 已提交
614 615
                list_elem.append(elem)
        return list_character, list_elem
L
LDOUBLEV 已提交
616

M
MissPenguin 已提交
617 618 619 620 621
    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
L
LDOUBLEV 已提交
622

M
MissPenguin 已提交
623 624 625 626 627 628
    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
L
LDOUBLEV 已提交
629

M
MissPenguin 已提交
630 631 632 633 634 635 636 637
    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]
L
LDOUBLEV 已提交
638 639
        structure = structure + [0] * (self.max_elem_length + 2 - len(structure)
                                       )
M
MissPenguin 已提交
640 641 642 643 644
        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()
L
LDOUBLEV 已提交
645 646
        td_idx_list = np.logical_or(structure == elem_char_idx1,
                                    structure == elem_char_idx2)
M
MissPenguin 已提交
647
        td_idx_list = np.where(td_idx_list)[0]
L
LDOUBLEV 已提交
648 649 650

        structure_mask = np.ones(
            (self.max_elem_length + 2, 1), dtype=np.float32)
M
MissPenguin 已提交
651
        bbox_list = np.zeros((self.max_elem_length + 2, 4), dtype=np.float32)
L
LDOUBLEV 已提交
652 653
        bbox_list_mask = np.zeros(
            (self.max_elem_length + 2, 1), dtype=np.float32)
M
MissPenguin 已提交
654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679
        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')
L
LDOUBLEV 已提交
680 681 682 683 684
        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
        ])
M
MissPenguin 已提交
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 732 733 734 735
        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" \
L
LDOUBLEV 已提交
736
                              % beg_or_end
M
MissPenguin 已提交
737 738
        else:
            assert False, "Unsupport type %s in char_or_elem" \
L
LDOUBLEV 已提交
739
                              % char_or_elem
M
MissPenguin 已提交
740
        return idx
A
andyjpaddle 已提交
741 742 743 744 745 746 747 748 749 750


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 已提交
751 752
        super(SARLabelEncode, self).__init__(
            max_text_length, character_dict_path, use_space_char)
A
andyjpaddle 已提交
753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777

    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 已提交
778

A
andyjpaddle 已提交
779 780 781 782 783 784
        padded_text[:len(target)] = target
        data['label'] = np.array(padded_text)
        return data

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