magic_model.py 41.2 KB
Newer Older
icecraft's avatar
icecraft committed
1
import enum
kernel.h@qq.com's avatar
kernel.h@qq.com committed
2

3
4
from magic_pdf.config.model_block_type import ModelBlockTypeEnum
from magic_pdf.config.ocr_content_type import CategoryId, ContentType
5
from magic_pdf.data.dataset import Dataset
6
from magic_pdf.libs.boxbase import (_is_in, _is_part_overlap, bbox_distance,
7
8
9
                                    bbox_relative_pos, box_area, calculate_iou,
                                    calculate_overlap_area_in_bbox1_area_ratio,
                                    get_overlap_area)
liukaiwen's avatar
liukaiwen committed
10
from magic_pdf.libs.coordinate_transform import get_scale_ratio
11
from magic_pdf.libs.local_math import float_gt
12
from magic_pdf.pre_proc.remove_bbox_overlap import _remove_overlap_between_bbox
liukaiwen's avatar
liukaiwen committed
13

14
CAPATION_OVERLAP_AREA_RATIO = 0.6
15
MERGE_BOX_OVERLAP_AREA_RATIO = 1.1
liukaiwen's avatar
liukaiwen committed
16

许瑞's avatar
许瑞 committed
17

icecraft's avatar
icecraft committed
18
19
20
21
22
23
24
25
class PosRelationEnum(enum.Enum):
    LEFT = 'left'
    RIGHT = 'right'
    UP = 'up'
    BOTTOM = 'bottom'
    ALL = 'all'


liukaiwen's avatar
liukaiwen committed
26
class MagicModel:
27
    """每个函数没有得到元素的时候返回空list."""
liukaiwen's avatar
liukaiwen committed
28
29
30

    def __fix_axis(self):
        for model_page_info in self.__model_list:
31
            need_remove_list = []
32
            page_no = model_page_info['page_info']['page_no']
liukaiwen's avatar
liukaiwen committed
33
            horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
34
                model_page_info, self.__docs.get_page(page_no)
liukaiwen's avatar
liukaiwen committed
35
            )
36
            layout_dets = model_page_info['layout_dets']
liukaiwen's avatar
liukaiwen committed
37
            for layout_det in layout_dets:
38

39
                if layout_det.get('bbox') is not None:
40
                    # 兼容直接输出bbox的模型数据,如paddle
41
                    x0, y0, x1, y1 = layout_det['bbox']
42
43
                else:
                    # 兼容直接输出poly的模型数据,如xxx
44
                    x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
45

liukaiwen's avatar
liukaiwen committed
46
47
48
49
50
51
                bbox = [
                    int(x0 / horizontal_scale_ratio),
                    int(y0 / vertical_scale_ratio),
                    int(x1 / horizontal_scale_ratio),
                    int(y1 / vertical_scale_ratio),
                ]
52
                layout_det['bbox'] = bbox
53
54
                # 删除高度或者宽度小于等于0的spans
                if bbox[2] - bbox[0] <= 0 or bbox[3] - bbox[1] <= 0:
liukaiwen's avatar
liukaiwen committed
55
56
57
58
                    need_remove_list.append(layout_det)
            for need_remove in need_remove_list:
                layout_dets.remove(need_remove)

59
    def __fix_by_remove_low_confidence(self):
60
61
        for model_page_info in self.__model_list:
            need_remove_list = []
62
            layout_dets = model_page_info['layout_dets']
63
            for layout_det in layout_dets:
64
                if layout_det['score'] <= 0.05:
65
66
67
68
69
                    need_remove_list.append(layout_det)
                else:
                    continue
            for need_remove in need_remove_list:
                layout_dets.remove(need_remove)
liukaiwen's avatar
liukaiwen committed
70

71
72
73
    def __fix_by_remove_high_iou_and_low_confidence(self):
        for model_page_info in self.__model_list:
            need_remove_list = []
74
            layout_dets = model_page_info['layout_dets']
75
76
77
78
            for layout_det1 in layout_dets:
                for layout_det2 in layout_dets:
                    if layout_det1 == layout_det2:
                        continue
79
                    if layout_det1['category_id'] in [
blue's avatar
blue committed
80
81
82
83
84
85
86
87
88
89
                        0,
                        1,
                        2,
                        3,
                        4,
                        5,
                        6,
                        7,
                        8,
                        9,
90
                    ] and layout_det2['category_id'] in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
blue's avatar
blue committed
91
                        if (
92
                            calculate_iou(layout_det1['bbox'], layout_det2['bbox'])
blue's avatar
blue committed
93
94
                            > 0.9
                        ):
95
                            if layout_det1['score'] < layout_det2['score']:
96
97
98
99
100
101
102
103
104
105
106
107
108
                                layout_det_need_remove = layout_det1
                            else:
                                layout_det_need_remove = layout_det2

                            if layout_det_need_remove not in need_remove_list:
                                need_remove_list.append(layout_det_need_remove)
                        else:
                            continue
                    else:
                        continue
            for need_remove in need_remove_list:
                layout_dets.remove(need_remove)

109
    def __init__(self, model_list: list, docs: Dataset):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
110
        self.__model_list = model_list
liukaiwen's avatar
liukaiwen committed
111
        self.__docs = docs
blue's avatar
blue committed
112
        """为所有模型数据添加bbox信息(缩放,poly->bbox)"""
kernel.h@qq.com's avatar
kernel.h@qq.com committed
113
        self.__fix_axis()
blue's avatar
blue committed
114
        """删除置信度特别低的模型数据(<0.05),提高质量"""
115
        self.__fix_by_remove_low_confidence()
blue's avatar
blue committed
116
        """删除高iou(>0.9)数据中置信度较低的那个"""
117
        self.__fix_by_remove_high_iou_and_low_confidence()
118
119
        self.__fix_footnote()

120
121
122
123
124
125
126
127
128
    def _bbox_distance(self, bbox1, bbox2):
        left, right, bottom, top = bbox_relative_pos(bbox1, bbox2)
        flags = [left, right, bottom, top]
        count = sum([1 if v else 0 for v in flags])
        if count > 1:
            return float('inf')
        if left or right:
            l1 = bbox1[3] - bbox1[1]
            l2 = bbox2[3] - bbox2[1]
129
        else:
130
131
            l1 = bbox1[2] - bbox1[0]
            l2 = bbox2[2] - bbox2[0]
132

133
        if l2 > l1 and (l2 - l1) / l1 > 0.3:
134
135
            return float('inf')

136
137
        return bbox_distance(bbox1, bbox2)

138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
    def __fix_footnote(self):
        # 3: figure, 5: table, 7: footnote
        for model_page_info in self.__model_list:
            footnotes = []
            figures = []
            tables = []

            for obj in model_page_info['layout_dets']:
                if obj['category_id'] == 7:
                    footnotes.append(obj)
                elif obj['category_id'] == 3:
                    figures.append(obj)
                elif obj['category_id'] == 5:
                    tables.append(obj)
                if len(footnotes) * len(figures) == 0:
                    continue
154
155
156
157
158
159
160
161
162
163
164
165
            dis_figure_footnote = {}
            dis_table_footnote = {}

            for i in range(len(footnotes)):
                for j in range(len(figures)):
                    pos_flag_count = sum(
                        list(
                            map(
                                lambda x: 1 if x else 0,
                                bbox_relative_pos(
                                    footnotes[i]['bbox'], figures[j]['bbox']
                                ),
166
167
                            )
                        )
168
169
170
171
                    )
                    if pos_flag_count > 1:
                        continue
                    dis_figure_footnote[i] = min(
172
                        self._bbox_distance(figures[j]['bbox'], footnotes[i]['bbox']),
173
174
175
176
177
178
179
180
181
182
183
                        dis_figure_footnote.get(i, float('inf')),
                    )
            for i in range(len(footnotes)):
                for j in range(len(tables)):
                    pos_flag_count = sum(
                        list(
                            map(
                                lambda x: 1 if x else 0,
                                bbox_relative_pos(
                                    footnotes[i]['bbox'], tables[j]['bbox']
                                ),
184
185
                            )
                        )
186
187
188
                    )
                    if pos_flag_count > 1:
                        continue
189

190
                    dis_table_footnote[i] = min(
191
                        self._bbox_distance(tables[j]['bbox'], footnotes[i]['bbox']),
192
193
194
195
196
197
198
                        dis_table_footnote.get(i, float('inf')),
                    )
            for i in range(len(footnotes)):
                if i not in dis_figure_footnote:
                    continue
                if dis_table_footnote.get(i, float('inf')) > dis_figure_footnote[i]:
                    footnotes[i]['category_id'] = CategoryId.ImageFootnote
liukaiwen's avatar
liukaiwen committed
199
200
201
202
203
204
205
206

    def __reduct_overlap(self, bboxes):
        N = len(bboxes)
        keep = [True] * N
        for i in range(N):
            for j in range(N):
                if i == j:
                    continue
207
                if _is_in(bboxes[i]['bbox'], bboxes[j]['bbox']):
liukaiwen's avatar
liukaiwen committed
208
209
210
211
                    keep[i] = False
        return [bboxes[i] for i in range(N) if keep[i]]

    def __tie_up_category_by_distance(
212
        self, page_no, subject_category_id, object_category_id
liukaiwen's avatar
liukaiwen committed
213
    ):
214
215
        """假定每个 subject 最多有一个 object (可以有多个相邻的 object 合并为单个 object),每个 object
        只能属于一个 subject."""
liukaiwen's avatar
liukaiwen committed
216
        ret = []
217
        MAX_DIS_OF_POINT = 10**9 + 7
218
219
220
221
222
        """
        subject 和 object 的 bbox 会合并成一个大的 bbox (named: merged bbox)。
        筛选出所有和 merged bbox 有 overlap 且 overlap 面积大于 object 的面积的 subjects。
        再求出筛选出的 subjects 和 object 的最短距离
        """
223
224

        def search_overlap_between_boxes(subject_idx, object_idx):
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
            idxes = [subject_idx, object_idx]
            x0s = [all_bboxes[idx]['bbox'][0] for idx in idxes]
            y0s = [all_bboxes[idx]['bbox'][1] for idx in idxes]
            x1s = [all_bboxes[idx]['bbox'][2] for idx in idxes]
            y1s = [all_bboxes[idx]['bbox'][3] for idx in idxes]

            merged_bbox = [
                min(x0s),
                min(y0s),
                max(x1s),
                max(y1s),
            ]
            ratio = 0

            other_objects = list(
                map(
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
                    filter(
                        lambda x: x['category_id']
                        not in (object_category_id, subject_category_id),
                        self.__model_list[page_no]['layout_dets'],
                    ),
                )
            )
            for other_object in other_objects:
                ratio = max(
                    ratio,
252
253
254
                    get_overlap_area(merged_bbox, other_object['bbox'])
                    * 1.0
                    / box_area(all_bboxes[object_idx]['bbox']),
255
256
257
258
259
                )
                if ratio >= MERGE_BOX_OVERLAP_AREA_RATIO:
                    break

            return ratio
liukaiwen's avatar
liukaiwen committed
260

许瑞's avatar
许瑞 committed
261
        def may_find_other_nearest_bbox(subject_idx, object_idx):
262
            ret = float('inf')
263

许瑞's avatar
许瑞 committed
264
            x0 = min(
265
                all_bboxes[subject_idx]['bbox'][0], all_bboxes[object_idx]['bbox'][0]
许瑞's avatar
许瑞 committed
266
267
            )
            y0 = min(
268
                all_bboxes[subject_idx]['bbox'][1], all_bboxes[object_idx]['bbox'][1]
许瑞's avatar
许瑞 committed
269
270
            )
            x1 = max(
271
                all_bboxes[subject_idx]['bbox'][2], all_bboxes[object_idx]['bbox'][2]
许瑞's avatar
许瑞 committed
272
273
            )
            y1 = max(
274
                all_bboxes[subject_idx]['bbox'][3], all_bboxes[object_idx]['bbox'][3]
许瑞's avatar
许瑞 committed
275
            )
许瑞's avatar
许瑞 committed
276

许瑞's avatar
许瑞 committed
277
            object_area = abs(
278
                all_bboxes[object_idx]['bbox'][2] - all_bboxes[object_idx]['bbox'][0]
许瑞's avatar
许瑞 committed
279
            ) * abs(
280
                all_bboxes[object_idx]['bbox'][3] - all_bboxes[object_idx]['bbox'][1]
许瑞's avatar
许瑞 committed
281
            )
许瑞's avatar
许瑞 committed
282
283

            for i in range(len(all_bboxes)):
许瑞's avatar
许瑞 committed
284
285
                if (
                    i == subject_idx
286
                    or all_bboxes[i]['category_id'] != subject_category_id
许瑞's avatar
许瑞 committed
287
                ):
许瑞's avatar
许瑞 committed
288
                    continue
289
290
                if _is_part_overlap([x0, y0, x1, y1], all_bboxes[i]['bbox']) or _is_in(
                    all_bboxes[i]['bbox'], [x0, y0, x1, y1]
许瑞's avatar
许瑞 committed
291
                ):
292

许瑞's avatar
许瑞 committed
293
                    i_area = abs(
294
295
                        all_bboxes[i]['bbox'][2] - all_bboxes[i]['bbox'][0]
                    ) * abs(all_bboxes[i]['bbox'][3] - all_bboxes[i]['bbox'][1])
许瑞's avatar
许瑞 committed
296
                    if i_area >= object_area:
297
                        ret = min(float('inf'), dis[i][object_idx])
298

许瑞's avatar
许瑞 committed
299
300
            return ret

blue's avatar
blue committed
301
        def expand_bbbox(idxes):
302
303
304
305
            x0s = [all_bboxes[idx]['bbox'][0] for idx in idxes]
            y0s = [all_bboxes[idx]['bbox'][1] for idx in idxes]
            x1s = [all_bboxes[idx]['bbox'][2] for idx in idxes]
            y1s = [all_bboxes[idx]['bbox'][3] for idx in idxes]
blue's avatar
blue committed
306
307
            return min(x0s), min(y0s), max(x1s), max(y1s)

liukaiwen's avatar
liukaiwen committed
308
309
310
        subjects = self.__reduct_overlap(
            list(
                map(
311
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
liukaiwen's avatar
liukaiwen committed
312
                    filter(
313
314
                        lambda x: x['category_id'] == subject_category_id,
                        self.__model_list[page_no]['layout_dets'],
liukaiwen's avatar
liukaiwen committed
315
316
317
318
319
320
321
322
                    ),
                )
            )
        )

        objects = self.__reduct_overlap(
            list(
                map(
323
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
liukaiwen's avatar
liukaiwen committed
324
                    filter(
325
326
                        lambda x: x['category_id'] == object_category_id,
                        self.__model_list[page_no]['layout_dets'],
liukaiwen's avatar
liukaiwen committed
327
328
329
330
331
332
                    ),
                )
            )
        )
        subject_object_relation_map = {}

许瑞's avatar
许瑞 committed
333
        subjects.sort(
334
            key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2
许瑞's avatar
许瑞 committed
335
        )  # get the distance !
liukaiwen's avatar
liukaiwen committed
336
337
338
339

        all_bboxes = []

        for v in subjects:
许瑞's avatar
许瑞 committed
340
341
            all_bboxes.append(
                {
342
343
344
                    'category_id': subject_category_id,
                    'bbox': v['bbox'],
                    'score': v['score'],
许瑞's avatar
许瑞 committed
345
346
                }
            )
liukaiwen's avatar
liukaiwen committed
347
348

        for v in objects:
许瑞's avatar
许瑞 committed
349
350
            all_bboxes.append(
                {
351
352
353
                    'category_id': object_category_id,
                    'bbox': v['bbox'],
                    'score': v['score'],
许瑞's avatar
许瑞 committed
354
355
                }
            )
liukaiwen's avatar
liukaiwen committed
356
357
358
359
360
361
362

        N = len(all_bboxes)
        dis = [[MAX_DIS_OF_POINT] * N for _ in range(N)]

        for i in range(N):
            for j in range(i):
                if (
363
364
                    all_bboxes[i]['category_id'] == subject_category_id
                    and all_bboxes[j]['category_id'] == subject_category_id
liukaiwen's avatar
liukaiwen committed
365
366
367
                ):
                    continue

368
369
370
371
                subject_idx, object_idx = i, j
                if all_bboxes[j]['category_id'] == subject_category_id:
                    subject_idx, object_idx = j, i

372
373
374
375
                if (
                    search_overlap_between_boxes(subject_idx, object_idx)
                    >= MERGE_BOX_OVERLAP_AREA_RATIO
                ):
376
377
378
379
                    dis[i][j] = float('inf')
                    dis[j][i] = dis[i][j]
                    continue

380
381
382
                dis[i][j] = self._bbox_distance(
                    all_bboxes[subject_idx]['bbox'], all_bboxes[object_idx]['bbox']
                )
liukaiwen's avatar
liukaiwen committed
383
384
385
386
387
                dis[j][i] = dis[i][j]

        used = set()
        for i in range(N):
            # 求第 i 个 subject 所关联的 object
388
            if all_bboxes[i]['category_id'] != subject_category_id:
liukaiwen's avatar
liukaiwen committed
389
390
391
392
393
394
395
396
397
398
399
                continue
            seen = set()
            candidates = []
            arr = []
            for j in range(N):

                pos_flag_count = sum(
                    list(
                        map(
                            lambda x: 1 if x else 0,
                            bbox_relative_pos(
400
                                all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
liukaiwen's avatar
liukaiwen committed
401
402
403
404
405
406
407
                            ),
                        )
                    )
                )
                if pos_flag_count > 1:
                    continue
                if (
408
                    all_bboxes[j]['category_id'] != object_category_id
409
410
                    or j in used
                    or dis[i][j] == MAX_DIS_OF_POINT
liukaiwen's avatar
liukaiwen committed
411
412
                ):
                    continue
blue's avatar
blue committed
413
                left, right, _, _ = bbox_relative_pos(
414
                    all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
blue's avatar
blue committed
415
                )  # 由  pos_flag_count 相关逻辑保证本段逻辑准确性
许瑞's avatar
许瑞 committed
416
                if left or right:
417
                    one_way_dis = all_bboxes[i]['bbox'][2] - all_bboxes[i]['bbox'][0]
许瑞's avatar
许瑞 committed
418
                else:
419
                    one_way_dis = all_bboxes[i]['bbox'][3] - all_bboxes[i]['bbox'][1]
许瑞's avatar
许瑞 committed
420
421
                if dis[i][j] > one_way_dis:
                    continue
liukaiwen's avatar
liukaiwen committed
422
423
424
425
                arr.append((dis[i][j], j))

            arr.sort(key=lambda x: x[0])
            if len(arr) > 0:
426
427
428
429
                """
                bug: 离该subject 最近的 object 可能跨越了其它的 subject。
                比如 [this subect] [some sbuject] [the nearest object of subject]
                """
许瑞's avatar
许瑞 committed
430
                if may_find_other_nearest_bbox(i, arr[0][1]) >= arr[0][0]:
许瑞's avatar
许瑞 committed
431

许瑞's avatar
许瑞 committed
432
433
                    candidates.append(arr[0][1])
                    seen.add(arr[0][1])
liukaiwen's avatar
liukaiwen committed
434
435
436
437
438
439
440
441
442
443

            # 已经获取初始种子
            for j in set(candidates):
                tmp = []
                for k in range(i + 1, N):
                    pos_flag_count = sum(
                        list(
                            map(
                                lambda x: 1 if x else 0,
                                bbox_relative_pos(
444
                                    all_bboxes[j]['bbox'], all_bboxes[k]['bbox']
liukaiwen's avatar
liukaiwen committed
445
446
447
448
449
450
451
452
453
                                ),
                            )
                        )
                    )

                    if pos_flag_count > 1:
                        continue

                    if (
454
                        all_bboxes[k]['category_id'] != object_category_id
455
456
457
                        or k in used
                        or k in seen
                        or dis[j][k] == MAX_DIS_OF_POINT
许瑞's avatar
许瑞 committed
458
                        or dis[j][k] > dis[i][j]
liukaiwen's avatar
liukaiwen committed
459
460
                    ):
                        continue
许瑞's avatar
许瑞 committed
461

liukaiwen's avatar
liukaiwen committed
462
                    is_nearest = True
463
464
                    for ni in range(i + 1, N):
                        if ni in (j, k) or ni in used or ni in seen:
liukaiwen's avatar
liukaiwen committed
465
466
                            continue

467
                        if not float_gt(dis[ni][k], dis[j][k]):
liukaiwen's avatar
liukaiwen committed
468
469
470
471
                            is_nearest = False
                            break

                    if is_nearest:
blue's avatar
blue committed
472
                        nx0, ny0, nx1, ny1 = expand_bbbox(list(seen) + [k])
473
                        n_dis = bbox_distance(
474
475
                            all_bboxes[i]['bbox'], [nx0, ny0, nx1, ny1]
                        )
blue's avatar
blue committed
476
477
                        if float_gt(dis[i][j], n_dis):
                            continue
liukaiwen's avatar
liukaiwen committed
478
479
480
481
482
483
484
485
486
                        tmp.append(k)
                        seen.add(k)

                candidates = tmp
                if len(candidates) == 0:
                    break

            # 已经获取到某个 figure 下所有的最靠近的 captions,以及最靠近这些 captions 的 captions 。
            # 先扩一下 bbox,
blue's avatar
blue committed
487
            ox0, oy0, ox1, oy1 = expand_bbbox(list(seen) + [i])
488
            ix0, iy0, ix1, iy1 = all_bboxes[i]['bbox']
liukaiwen's avatar
liukaiwen committed
489
490
491
492
493
494
495
496
497
498
499
500
501

            # 分成了 4 个截取空间,需要计算落在每个截取空间下 objects 合并后占据的矩形面积
            caption_poses = [
                [ox0, oy0, ix0, oy1],
                [ox0, oy0, ox1, iy0],
                [ox0, iy1, ox1, oy1],
                [ix1, oy0, ox1, oy1],
            ]

            caption_areas = []
            for bbox in caption_poses:
                embed_arr = []
                for idx in seen:
许瑞's avatar
许瑞 committed
502
503
                    if (
                        calculate_overlap_area_in_bbox1_area_ratio(
504
                            all_bboxes[idx]['bbox'], bbox
许瑞's avatar
许瑞 committed
505
506
507
                        )
                        > CAPATION_OVERLAP_AREA_RATIO
                    ):
liukaiwen's avatar
liukaiwen committed
508
509
510
                        embed_arr.append(idx)

                if len(embed_arr) > 0:
511
512
513
514
                    embed_x0 = min([all_bboxes[idx]['bbox'][0] for idx in embed_arr])
                    embed_y0 = min([all_bboxes[idx]['bbox'][1] for idx in embed_arr])
                    embed_x1 = max([all_bboxes[idx]['bbox'][2] for idx in embed_arr])
                    embed_y1 = max([all_bboxes[idx]['bbox'][3] for idx in embed_arr])
liukaiwen's avatar
liukaiwen committed
515
516
517
518
519
520
521
522
523
524
525
526
                    caption_areas.append(
                        int(abs(embed_x1 - embed_x0) * abs(embed_y1 - embed_y0))
                    )
                else:
                    caption_areas.append(0)

            subject_object_relation_map[i] = []
            if max(caption_areas) > 0:
                max_area_idx = caption_areas.index(max(caption_areas))
                caption_bbox = caption_poses[max_area_idx]

                for j in seen:
许瑞's avatar
许瑞 committed
527
528
                    if (
                        calculate_overlap_area_in_bbox1_area_ratio(
529
                            all_bboxes[j]['bbox'], caption_bbox
许瑞's avatar
许瑞 committed
530
531
532
                        )
                        > CAPATION_OVERLAP_AREA_RATIO
                    ):
liukaiwen's avatar
liukaiwen committed
533
534
535
536
537
                        used.add(j)
                        subject_object_relation_map[i].append(j)

        for i in sorted(subject_object_relation_map.keys()):
            result = {
538
539
540
                'subject_body': all_bboxes[i]['bbox'],
                'all': all_bboxes[i]['bbox'],
                'score': all_bboxes[i]['score'],
liukaiwen's avatar
liukaiwen committed
541
542
543
544
            }

            if len(subject_object_relation_map[i]) > 0:
                x0 = min(
545
                    [all_bboxes[j]['bbox'][0] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
546
547
                )
                y0 = min(
548
                    [all_bboxes[j]['bbox'][1] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
549
550
                )
                x1 = max(
551
                    [all_bboxes[j]['bbox'][2] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
552
553
                )
                y1 = max(
554
                    [all_bboxes[j]['bbox'][3] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
555
                )
556
557
558
559
560
561
                result['object_body'] = [x0, y0, x1, y1]
                result['all'] = [
                    min(x0, all_bboxes[i]['bbox'][0]),
                    min(y0, all_bboxes[i]['bbox'][1]),
                    max(x1, all_bboxes[i]['bbox'][2]),
                    max(y1, all_bboxes[i]['bbox'][3]),
liukaiwen's avatar
liukaiwen committed
562
563
564
565
566
567
568
                ]
            ret.append(result)

        total_subject_object_dis = 0
        # 计算已经配对的 distance 距离
        for i in subject_object_relation_map.keys():
            for j in subject_object_relation_map[i]:
569
                total_subject_object_dis += bbox_distance(
570
                    all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
liukaiwen's avatar
liukaiwen committed
571
572
573
574
575
576
577
578
579
580
581
                )

        # 计算未匹配的 subject 和 object 的距离(非精确版)
        with_caption_subject = set(
            [
                key
                for key in subject_object_relation_map.keys()
                if len(subject_object_relation_map[i]) > 0
            ]
        )
        for i in range(N):
582
            if all_bboxes[i]['category_id'] != object_category_id or i in used:
liukaiwen's avatar
liukaiwen committed
583
584
585
586
                continue
            candidates = []
            for j in range(N):
                if (
587
                    all_bboxes[j]['category_id'] != subject_category_id
588
                    or j in with_caption_subject
liukaiwen's avatar
liukaiwen committed
589
590
591
592
593
594
595
596
597
                ):
                    continue
                candidates.append((dis[i][j], j))
            if len(candidates) > 0:
                candidates.sort(key=lambda x: x[0])
                total_subject_object_dis += candidates[0][1]
                with_caption_subject.add(j)
        return ret, total_subject_object_dis

598
    def __tie_up_category_by_distance_v2(
icecraft's avatar
icecraft committed
599
600
601
602
603
        self,
        page_no: int,
        subject_category_id: int,
        object_category_id: int,
        priority_pos: PosRelationEnum,
604
    ):
icecraft's avatar
icecraft committed
605
        """_summary_
606

icecraft's avatar
icecraft committed
607
608
609
610
611
612
613
614
615
        Args:
            page_no (int): _description_
            subject_category_id (int): _description_
            object_category_id (int): _description_
            priority_pos (PosRelationEnum): _description_

        Returns:
            _type_: _description_
        """
icecraft's avatar
icecraft committed
616
        AXIS_MULPLICITY = 0.5
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
        subjects = self.__reduct_overlap(
            list(
                map(
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
                    filter(
                        lambda x: x['category_id'] == subject_category_id,
                        self.__model_list[page_no]['layout_dets'],
                    ),
                )
            )
        )

        objects = self.__reduct_overlap(
            list(
                map(
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
                    filter(
                        lambda x: x['category_id'] == object_category_id,
                        self.__model_list[page_no]['layout_dets'],
                    ),
                )
            )
        )
640
        M = len(objects)
641
642
643

        subjects.sort(key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2)
        objects.sort(key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2)
644
645
646
647
648
649
650
651
652
653

        sub_obj_map_h = {i: [] for i in range(len(subjects))}

        dis_by_directions = {
            'top': [[-1, float('inf')]] * M,
            'bottom': [[-1, float('inf')]] * M,
            'left': [[-1, float('inf')]] * M,
            'right': [[-1, float('inf')]] * M,
        }

654
        for i, obj in enumerate(objects):
655
656
657
658
659
            l_x_axis, l_y_axis = (
                obj['bbox'][2] - obj['bbox'][0],
                obj['bbox'][3] - obj['bbox'][1],
            )
            axis_unit = min(l_x_axis, l_y_axis)
660
661
            for j, sub in enumerate(subjects):

icecraft's avatar
icecraft committed
662
663
                bbox1, bbox2, _ = _remove_overlap_between_bbox(
                    objects[i]['bbox'], subjects[j]['bbox']
664
                )
icecraft's avatar
icecraft committed
665
                left, right, bottom, top = bbox_relative_pos(bbox1, bbox2)
666
667
                flags = [left, right, bottom, top]
                if sum([1 if v else 0 for v in flags]) > 1:
668
669
                    continue

670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
                if left:
                    if dis_by_directions['left'][i][1] > bbox_distance(
                        obj['bbox'], sub['bbox']
                    ):
                        dis_by_directions['left'][i] = [
                            j,
                            bbox_distance(obj['bbox'], sub['bbox']),
                        ]
                if right:
                    if dis_by_directions['right'][i][1] > bbox_distance(
                        obj['bbox'], sub['bbox']
                    ):
                        dis_by_directions['right'][i] = [
                            j,
                            bbox_distance(obj['bbox'], sub['bbox']),
                        ]
                if bottom:
                    if dis_by_directions['bottom'][i][1] > bbox_distance(
                        obj['bbox'], sub['bbox']
                    ):
                        dis_by_directions['bottom'][i] = [
                            j,
                            bbox_distance(obj['bbox'], sub['bbox']),
                        ]
                if top:
                    if dis_by_directions['top'][i][1] > bbox_distance(
                        obj['bbox'], sub['bbox']
                    ):
                        dis_by_directions['top'][i] = [
                            j,
                            bbox_distance(obj['bbox'], sub['bbox']),
                        ]
icecraft's avatar
icecraft committed
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722

            if (
                dis_by_directions['top'][i][1] != float('inf')
                and dis_by_directions['bottom'][i][1] != float('inf')
                and priority_pos in (PosRelationEnum.BOTTOM, PosRelationEnum.UP)
            ):
                RATIO = 3
                if (
                    abs(
                        dis_by_directions['top'][i][1]
                        - dis_by_directions['bottom'][i][1]
                    )
                    < RATIO * axis_unit
                ):

                    if priority_pos == PosRelationEnum.BOTTOM:
                        sub_obj_map_h[dis_by_directions['bottom'][i][0]].append(i)
                    else:
                        sub_obj_map_h[dis_by_directions['top'][i][0]].append(i)
                    continue

723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
            if dis_by_directions['left'][i][1] != float('inf') or dis_by_directions[
                'right'
            ][i][1] != float('inf'):
                if dis_by_directions['left'][i][1] != float(
                    'inf'
                ) and dis_by_directions['right'][i][1] != float('inf'):
                    if AXIS_MULPLICITY * axis_unit >= abs(
                        dis_by_directions['left'][i][1]
                        - dis_by_directions['right'][i][1]
                    ):
                        left_sub_bbox = subjects[dis_by_directions['left'][i][0]][
                            'bbox'
                        ]
                        right_sub_bbox = subjects[dis_by_directions['right'][i][0]][
                            'bbox'
                        ]

                        left_sub_bbox_y_axis = left_sub_bbox[3] - left_sub_bbox[1]
                        right_sub_bbox_y_axis = right_sub_bbox[3] - right_sub_bbox[1]

icecraft's avatar
icecraft committed
743
744
745
746
747
                        if (
                            abs(left_sub_bbox_y_axis - l_y_axis)
                            + dis_by_directions['left'][i][0]
                            > abs(right_sub_bbox_y_axis - l_y_axis)
                            + dis_by_directions['right'][i][0]
748
749
750
751
752
753
                        ):
                            left_or_right = dis_by_directions['right'][i]
                        else:
                            left_or_right = dis_by_directions['left'][i]
                    else:
                        left_or_right = dis_by_directions['left'][i]
icecraft's avatar
icecraft committed
754
                        if left_or_right[1] > dis_by_directions['right'][i][1]:
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
                            left_or_right = dis_by_directions['right'][i]
                else:
                    left_or_right = dis_by_directions['left'][i]
                    if left_or_right[1] == float('inf'):
                        left_or_right = dis_by_directions['right'][i]
            else:
                left_or_right = [-1, float('inf')]

            if dis_by_directions['top'][i][1] != float('inf') or dis_by_directions[
                'bottom'
            ][i][1] != float('inf'):
                if dis_by_directions['top'][i][1] != float('inf') and dis_by_directions[
                    'bottom'
                ][i][1] != float('inf'):
                    if AXIS_MULPLICITY * axis_unit >= abs(
                        dis_by_directions['top'][i][1]
                        - dis_by_directions['bottom'][i][1]
                    ):
                        top_bottom = subjects[dis_by_directions['bottom'][i][0]]['bbox']
                        bottom_top = subjects[dis_by_directions['top'][i][0]]['bbox']

                        top_bottom_x_axis = top_bottom[2] - top_bottom[0]
                        bottom_top_x_axis = bottom_top[2] - bottom_top[0]
icecraft's avatar
icecraft committed
778
779
780
781
782
783
                        if (
                            abs(top_bottom_x_axis - l_x_axis)
                            + dis_by_directions['bottom'][i][1]
                            > abs(bottom_top_x_axis - l_x_axis)
                            + dis_by_directions['top'][i][1]
                        ):
784
                            top_or_bottom = dis_by_directions['top'][i]
icecraft's avatar
icecraft committed
785
786
                        else:
                            top_or_bottom = dis_by_directions['bottom'][i]
787
788
                    else:
                        top_or_bottom = dis_by_directions['top'][i]
icecraft's avatar
icecraft committed
789
                        if top_or_bottom[1] > dis_by_directions['bottom'][i][1]:
790
791
792
793
794
                            top_or_bottom = dis_by_directions['bottom'][i]
                else:
                    top_or_bottom = dis_by_directions['top'][i]
                    if top_or_bottom[1] == float('inf'):
                        top_or_bottom = dis_by_directions['bottom'][i]
795
            else:
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
                top_or_bottom = [-1, float('inf')]

            if left_or_right[1] != float('inf') or top_or_bottom[1] != float('inf'):
                if left_or_right[1] != float('inf') and top_or_bottom[1] != float(
                    'inf'
                ):
                    if AXIS_MULPLICITY * axis_unit >= abs(
                        left_or_right[1] - top_or_bottom[1]
                    ):
                        y_axis_bbox = subjects[left_or_right[0]]['bbox']
                        x_axis_bbox = subjects[top_or_bottom[0]]['bbox']

                        if (
                            abs((x_axis_bbox[2] - x_axis_bbox[0]) - l_x_axis) / l_x_axis
                            > abs((y_axis_bbox[3] - y_axis_bbox[1]) - l_y_axis)
                            / l_y_axis
                        ):
                            sub_obj_map_h[left_or_right[0]].append(i)
                        else:
                            sub_obj_map_h[top_or_bottom[0]].append(i)
                    else:
                        if left_or_right[1] > top_or_bottom[1]:
                            sub_obj_map_h[top_or_bottom[0]].append(i)
                        else:
                            sub_obj_map_h[left_or_right[0]].append(i)
                else:
                    if left_or_right[1] != float('inf'):
                        sub_obj_map_h[left_or_right[0]].append(i)
                    else:
                        sub_obj_map_h[top_or_bottom[0]].append(i)
826
827
828
829
        ret = []
        for i in sub_obj_map_h.keys():
            ret.append(
                {
icecraft's avatar
icecraft committed
830
831
832
833
                    'sub_bbox': {
                        'bbox': subjects[i]['bbox'],
                        'score': subjects[i]['score'],
                    },
834
835
836
837
                    'obj_bboxes': [
                        {'score': objects[j]['score'], 'bbox': objects[j]['bbox']}
                        for j in sub_obj_map_h[i]
                    ],
838
839
840
841
842
843
                    'sub_idx': i,
                }
            )
        return ret

    def get_imgs_v2(self, page_no: int):
icecraft's avatar
icecraft committed
844
845
846
        with_captions = self.__tie_up_category_by_distance_v2(
            page_no, 3, 4, PosRelationEnum.BOTTOM
        )
847
        with_footnotes = self.__tie_up_category_by_distance_v2(
icecraft's avatar
icecraft committed
848
            page_no, 3, CategoryId.ImageFootnote, PosRelationEnum.ALL
849
850
851
852
        )
        ret = []
        for v in with_captions:
            record = {
853
854
                'image_body': v['sub_bbox'],
                'image_caption_list': v['obj_bboxes'],
855
856
857
            }
            filter_idx = v['sub_idx']
            d = next(filter(lambda x: x['sub_idx'] == filter_idx, with_footnotes))
858
            record['image_footnote_list'] = d['obj_bboxes']
859
860
861
862
            ret.append(record)
        return ret

    def get_tables_v2(self, page_no: int) -> list:
icecraft's avatar
icecraft committed
863
864
865
866
867
868
        with_captions = self.__tie_up_category_by_distance_v2(
            page_no, 5, 6, PosRelationEnum.UP
        )
        with_footnotes = self.__tie_up_category_by_distance_v2(
            page_no, 5, 7, PosRelationEnum.ALL
        )
869
870
871
        ret = []
        for v in with_captions:
            record = {
872
873
                'table_body': v['sub_bbox'],
                'table_caption_list': v['obj_bboxes'],
874
875
876
            }
            filter_idx = v['sub_idx']
            d = next(filter(lambda x: x['sub_idx'] == filter_idx, with_footnotes))
877
            record['table_footnote_list'] = d['obj_bboxes']
878
879
880
            ret.append(record)
        return ret

blue's avatar
blue committed
881
    def get_imgs(self, page_no: int):
882
883
884
        with_captions, _ = self.__tie_up_category_by_distance(page_no, 3, 4)
        with_footnotes, _ = self.__tie_up_category_by_distance(
            page_no, 3, CategoryId.ImageFootnote
blue's avatar
blue committed
885
        )
886
887
888
889
890
891
892
893
894
        ret = []
        N, M = len(with_captions), len(with_footnotes)
        assert N == M
        for i in range(N):
            record = {
                'score': with_captions[i]['score'],
                'img_caption_bbox': with_captions[i].get('object_body', None),
                'img_body_bbox': with_captions[i]['subject_body'],
                'img_footnote_bbox': with_footnotes[i].get('object_body', None),
liukaiwen's avatar
liukaiwen committed
895
            }
896
897
898
899
900
901
902
903

            x0 = min(with_captions[i]['all'][0], with_footnotes[i]['all'][0])
            y0 = min(with_captions[i]['all'][1], with_footnotes[i]['all'][1])
            x1 = max(with_captions[i]['all'][2], with_footnotes[i]['all'][2])
            y1 = max(with_captions[i]['all'][3], with_footnotes[i]['all'][3])
            record['bbox'] = [x0, y0, x1, y1]
            ret.append(record)
        return ret
liukaiwen's avatar
liukaiwen committed
904
905

    def get_tables(
906
        self, page_no: int
liukaiwen's avatar
liukaiwen committed
907
908
909
910
911
912
913
914
    ) -> list:  # 3个坐标, caption, table主体,table-note
        with_captions, _ = self.__tie_up_category_by_distance(page_no, 5, 6)
        with_footnotes, _ = self.__tie_up_category_by_distance(page_no, 5, 7)
        ret = []
        N, M = len(with_captions), len(with_footnotes)
        assert N == M
        for i in range(N):
            record = {
915
916
917
918
                'score': with_captions[i]['score'],
                'table_caption_bbox': with_captions[i].get('object_body', None),
                'table_body_bbox': with_captions[i]['subject_body'],
                'table_footnote_bbox': with_footnotes[i].get('object_body', None),
liukaiwen's avatar
liukaiwen committed
919
920
            }

921
922
923
924
925
            x0 = min(with_captions[i]['all'][0], with_footnotes[i]['all'][0])
            y0 = min(with_captions[i]['all'][1], with_footnotes[i]['all'][1])
            x1 = max(with_captions[i]['all'][2], with_footnotes[i]['all'][2])
            y1 = max(with_captions[i]['all'][3], with_footnotes[i]['all'][3])
            record['bbox'] = [x0, y0, x1, y1]
liukaiwen's avatar
liukaiwen committed
926
927
928
929
            ret.append(record)
        return ret

    def get_equations(self, page_no: int) -> list:  # 有坐标,也有字
930
        inline_equations = self.__get_blocks_by_type(
931
            ModelBlockTypeEnum.EMBEDDING.value, page_no, ['latex']
932
933
        )
        interline_equations = self.__get_blocks_by_type(
934
            ModelBlockTypeEnum.ISOLATED.value, page_no, ['latex']
935
936
937
938
        )
        interline_equations_blocks = self.__get_blocks_by_type(
            ModelBlockTypeEnum.ISOLATE_FORMULA.value, page_no
        )
liukaiwen's avatar
liukaiwen committed
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
        return inline_equations, interline_equations, interline_equations_blocks

    def get_discarded(self, page_no: int) -> list:  # 自研模型,只有坐标
        blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.ABANDON.value, page_no)
        return blocks

    def get_text_blocks(self, page_no: int) -> list:  # 自研模型搞的,只有坐标,没有字
        blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.PLAIN_TEXT.value, page_no)
        return blocks

    def get_title_blocks(self, page_no: int) -> list:  # 自研模型,只有坐标,没字
        blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.TITLE.value, page_no)
        return blocks

    def get_ocr_text(self, page_no: int) -> list:  # paddle 搞的,有字也有坐标
        text_spans = []
        model_page_info = self.__model_list[page_no]
956
        layout_dets = model_page_info['layout_dets']
liukaiwen's avatar
liukaiwen committed
957
        for layout_det in layout_dets:
958
            if layout_det['category_id'] == '15':
liukaiwen's avatar
liukaiwen committed
959
                span = {
960
961
                    'bbox': layout_det['bbox'],
                    'content': layout_det['text'],
liukaiwen's avatar
liukaiwen committed
962
963
964
965
966
                }
                text_spans.append(span)
        return text_spans

    def get_all_spans(self, page_no: int) -> list:
967

968
969
970
971
972
973
        def remove_duplicate_spans(spans):
            new_spans = []
            for span in spans:
                if not any(span == existing_span for existing_span in new_spans):
                    new_spans.append(span)
            return new_spans
blue's avatar
blue committed
974

liukaiwen's avatar
liukaiwen committed
975
976
        all_spans = []
        model_page_info = self.__model_list[page_no]
977
        layout_dets = model_page_info['layout_dets']
liukaiwen's avatar
liukaiwen committed
978
979
980
        allow_category_id_list = [3, 5, 13, 14, 15]
        """当成span拼接的"""
        #  3: 'image', # 图片
981
        #  5: 'table',       # 表格
liukaiwen's avatar
liukaiwen committed
982
983
984
985
        #  13: 'inline_equation',     # 行内公式
        #  14: 'interline_equation',      # 行间公式
        #  15: 'text',      # ocr识别文本
        for layout_det in layout_dets:
986
            category_id = layout_det['category_id']
liukaiwen's avatar
liukaiwen committed
987
            if category_id in allow_category_id_list:
988
                span = {'bbox': layout_det['bbox'], 'score': layout_det['score']}
liukaiwen's avatar
liukaiwen committed
989
                if category_id == 3:
990
                    span['type'] = ContentType.Image
liukaiwen's avatar
liukaiwen committed
991
                elif category_id == 5:
992
                    # 获取table模型结果
993
994
                    latex = layout_det.get('latex', None)
                    html = layout_det.get('html', None)
995
                    if latex:
996
                        span['latex'] = latex
997
                    elif html:
998
999
                        span['html'] = html
                    span['type'] = ContentType.Table
liukaiwen's avatar
liukaiwen committed
1000
                elif category_id == 13:
1001
1002
                    span['content'] = layout_det['latex']
                    span['type'] = ContentType.InlineEquation
liukaiwen's avatar
liukaiwen committed
1003
                elif category_id == 14:
1004
1005
                    span['content'] = layout_det['latex']
                    span['type'] = ContentType.InterlineEquation
liukaiwen's avatar
liukaiwen committed
1006
                elif category_id == 15:
1007
1008
                    span['content'] = layout_det['text']
                    span['type'] = ContentType.Text
liukaiwen's avatar
liukaiwen committed
1009
                all_spans.append(span)
1010
        return remove_duplicate_spans(all_spans)
liukaiwen's avatar
liukaiwen committed
1011
1012
1013

    def get_page_size(self, page_no: int):  # 获取页面宽高
        # 获取当前页的page对象
1014
        page = self.__docs.get_page(page_no).get_page_info()
liukaiwen's avatar
liukaiwen committed
1015
        # 获取当前页的宽高
1016
1017
        page_w = page.w
        page_h = page.h
liukaiwen's avatar
liukaiwen committed
1018
1019
        return page_w, page_h

1020
1021
1022
    def __get_blocks_by_type(
        self, type: int, page_no: int, extra_col: list[str] = []
    ) -> list:
liukaiwen's avatar
liukaiwen committed
1023
1024
        blocks = []
        for page_dict in self.__model_list:
1025
1026
1027
            layout_dets = page_dict.get('layout_dets', [])
            page_info = page_dict.get('page_info', {})
            page_number = page_info.get('page_no', -1)
liukaiwen's avatar
liukaiwen committed
1028
1029
1030
            if page_no != page_number:
                continue
            for item in layout_dets:
1031
1032
                category_id = item.get('category_id', -1)
                bbox = item.get('bbox', None)
liukaiwen's avatar
liukaiwen committed
1033

liukaiwen's avatar
liukaiwen committed
1034
                if category_id == type:
1035
                    block = {
1036
1037
                        'bbox': bbox,
                        'score': item.get('score'),
1038
                    }
liukaiwen's avatar
liukaiwen committed
1039
1040
1041
1042
1043
                    for col in extra_col:
                        block[col] = item.get(col, None)
                    blocks.append(block)
        return blocks

许瑞's avatar
许瑞 committed
1044
1045
1046
    def get_model_list(self, page_no):
        return self.__model_list[page_no]