magic_model.py 28.2 KB
Newer Older
liukaiwen's avatar
liukaiwen committed
1
import json
kernel.h@qq.com's avatar
kernel.h@qq.com committed
2

3
4
5
6
from magic_pdf.libs.boxbase import (_is_in, _is_part_overlap, bbox_distance,
                                    bbox_relative_pos, calculate_iou,
                                    calculate_overlap_area_in_bbox1_area_ratio)
from magic_pdf.libs.commons import fitz, join_path
liukaiwen's avatar
liukaiwen committed
7
from magic_pdf.libs.coordinate_transform import get_scale_ratio
8
from magic_pdf.libs.local_math import float_gt
liukaiwen's avatar
liukaiwen committed
9
from magic_pdf.libs.ModelBlockTypeEnum import ModelBlockTypeEnum
10
11
12
from magic_pdf.libs.ocr_content_type import CategoryId, ContentType
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
liukaiwen's avatar
liukaiwen committed
13

14
CAPATION_OVERLAP_AREA_RATIO = 0.6
liukaiwen's avatar
liukaiwen committed
15

许瑞's avatar
许瑞 committed
16

liukaiwen's avatar
liukaiwen committed
17
class MagicModel:
18
    """每个函数没有得到元素的时候返回空list."""
liukaiwen's avatar
liukaiwen committed
19
20
21

    def __fix_axis(self):
        for model_page_info in self.__model_list:
22
            need_remove_list = []
23
            page_no = model_page_info['page_info']['page_no']
liukaiwen's avatar
liukaiwen committed
24
25
26
            horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
                model_page_info, self.__docs[page_no]
            )
27
            layout_dets = model_page_info['layout_dets']
liukaiwen's avatar
liukaiwen committed
28
            for layout_det in layout_dets:
29

30
                if layout_det.get('bbox') is not None:
31
                    # 兼容直接输出bbox的模型数据,如paddle
32
                    x0, y0, x1, y1 = layout_det['bbox']
33
34
                else:
                    # 兼容直接输出poly的模型数据,如xxx
35
                    x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
36

liukaiwen's avatar
liukaiwen committed
37
38
39
40
41
42
                bbox = [
                    int(x0 / horizontal_scale_ratio),
                    int(y0 / vertical_scale_ratio),
                    int(x1 / horizontal_scale_ratio),
                    int(y1 / vertical_scale_ratio),
                ]
43
                layout_det['bbox'] = bbox
44
45
                # 删除高度或者宽度小于等于0的spans
                if bbox[2] - bbox[0] <= 0 or bbox[3] - bbox[1] <= 0:
liukaiwen's avatar
liukaiwen committed
46
47
48
49
                    need_remove_list.append(layout_det)
            for need_remove in need_remove_list:
                layout_dets.remove(need_remove)

50
    def __fix_by_remove_low_confidence(self):
51
52
        for model_page_info in self.__model_list:
            need_remove_list = []
53
            layout_dets = model_page_info['layout_dets']
54
            for layout_det in layout_dets:
55
                if layout_det['score'] <= 0.05:
56
57
58
59
60
                    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
61

62
63
64
    def __fix_by_remove_high_iou_and_low_confidence(self):
        for model_page_info in self.__model_list:
            need_remove_list = []
65
            layout_dets = model_page_info['layout_dets']
66
67
68
69
            for layout_det1 in layout_dets:
                for layout_det2 in layout_dets:
                    if layout_det1 == layout_det2:
                        continue
70
                    if layout_det1['category_id'] in [
blue's avatar
blue committed
71
72
73
74
75
76
77
78
79
80
                        0,
                        1,
                        2,
                        3,
                        4,
                        5,
                        6,
                        7,
                        8,
                        9,
81
                    ] and layout_det2['category_id'] in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
blue's avatar
blue committed
82
                        if (
83
                            calculate_iou(layout_det1['bbox'], layout_det2['bbox'])
blue's avatar
blue committed
84
85
                            > 0.9
                        ):
86
                            if layout_det1['score'] < layout_det2['score']:
87
88
89
90
91
92
93
94
95
96
97
98
99
                                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)

liukaiwen's avatar
liukaiwen committed
100
    def __init__(self, model_list: list, docs: fitz.Document):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
101
        self.__model_list = model_list
liukaiwen's avatar
liukaiwen committed
102
        self.__docs = docs
blue's avatar
blue committed
103
        """为所有模型数据添加bbox信息(缩放,poly->bbox)"""
kernel.h@qq.com's avatar
kernel.h@qq.com committed
104
        self.__fix_axis()
blue's avatar
blue committed
105
        """删除置信度特别低的模型数据(<0.05),提高质量"""
106
        self.__fix_by_remove_low_confidence()
blue's avatar
blue committed
107
        """删除高iou(>0.9)数据中置信度较低的那个"""
108
        self.__fix_by_remove_high_iou_and_low_confidence()
109
110
111
112
113
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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
        self.__fix_footnote()

    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
                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']
                                    ),
                                )
                            )
                        )
                        if pos_flag_count > 1:
                            continue
                        dis_figure_footnote[i] = min(
                            bbox_distance(figures[j]['bbox'], footnotes[i]['bbox']),
                            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']
                                    ),
                                )
                            )
                        )
                        if pos_flag_count > 1:
                            continue

                        dis_table_footnote[i] = min(
                            bbox_distance(tables[j]['bbox'], footnotes[i]['bbox']),
                            dis_table_footnote.get(i, float('inf')),
                        )
                for i in range(len(footnotes)):
                    if dis_table_footnote.get(i, float('inf')) > dis_figure_footnote[i]:
                        footnotes[i]['category_id'] = CategoryId.ImageFootnote
liukaiwen's avatar
liukaiwen committed
170
171
172
173
174
175
176
177

    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
178
                if _is_in(bboxes[i]['bbox'], bboxes[j]['bbox']):
liukaiwen's avatar
liukaiwen committed
179
180
181
182
                    keep[i] = False
        return [bboxes[i] for i in range(N) if keep[i]]

    def __tie_up_category_by_distance(
183
        self, page_no, subject_category_id, object_category_id
liukaiwen's avatar
liukaiwen committed
184
    ):
185
186
        """假定每个 subject 最多有一个 object (可以有多个相邻的 object 合并为单个 object),每个 object
        只能属于一个 subject."""
liukaiwen's avatar
liukaiwen committed
187
        ret = []
188
        MAX_DIS_OF_POINT = 10**9 + 7
189
190
191
192
193
        """
        subject 和 object 的 bbox 会合并成一个大的 bbox (named: merged bbox)。
        筛选出所有和 merged bbox 有 overlap 且 overlap 面积大于 object 的面积的 subjects。
        再求出筛选出的 subjects 和 object 的最短距离
        """
liukaiwen's avatar
liukaiwen committed
194

许瑞's avatar
许瑞 committed
195
        def may_find_other_nearest_bbox(subject_idx, object_idx):
196
            ret = float('inf')
197

许瑞's avatar
许瑞 committed
198
            x0 = min(
199
                all_bboxes[subject_idx]['bbox'][0], all_bboxes[object_idx]['bbox'][0]
许瑞's avatar
许瑞 committed
200
201
            )
            y0 = min(
202
                all_bboxes[subject_idx]['bbox'][1], all_bboxes[object_idx]['bbox'][1]
许瑞's avatar
许瑞 committed
203
204
            )
            x1 = max(
205
                all_bboxes[subject_idx]['bbox'][2], all_bboxes[object_idx]['bbox'][2]
许瑞's avatar
许瑞 committed
206
207
            )
            y1 = max(
208
                all_bboxes[subject_idx]['bbox'][3], all_bboxes[object_idx]['bbox'][3]
许瑞's avatar
许瑞 committed
209
            )
许瑞's avatar
许瑞 committed
210

许瑞's avatar
许瑞 committed
211
            object_area = abs(
212
                all_bboxes[object_idx]['bbox'][2] - all_bboxes[object_idx]['bbox'][0]
许瑞's avatar
许瑞 committed
213
            ) * abs(
214
                all_bboxes[object_idx]['bbox'][3] - all_bboxes[object_idx]['bbox'][1]
许瑞's avatar
许瑞 committed
215
            )
许瑞's avatar
许瑞 committed
216
217

            for i in range(len(all_bboxes)):
许瑞's avatar
许瑞 committed
218
219
                if (
                    i == subject_idx
220
                    or all_bboxes[i]['category_id'] != subject_category_id
许瑞's avatar
许瑞 committed
221
                ):
许瑞's avatar
许瑞 committed
222
                    continue
223
224
                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
225
                ):
226

许瑞's avatar
许瑞 committed
227
                    i_area = abs(
228
229
                        all_bboxes[i]['bbox'][2] - all_bboxes[i]['bbox'][0]
                    ) * abs(all_bboxes[i]['bbox'][3] - all_bboxes[i]['bbox'][1])
许瑞's avatar
许瑞 committed
230
                    if i_area >= object_area:
231
                        ret = min(float('inf'), dis[i][object_idx])
232

许瑞's avatar
许瑞 committed
233
234
            return ret

blue's avatar
blue committed
235
        def expand_bbbox(idxes):
236
237
238
239
            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
240
241
            return min(x0s), min(y0s), max(x1s), max(y1s)

liukaiwen's avatar
liukaiwen committed
242
243
244
        subjects = self.__reduct_overlap(
            list(
                map(
245
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
liukaiwen's avatar
liukaiwen committed
246
                    filter(
247
248
                        lambda x: x['category_id'] == subject_category_id,
                        self.__model_list[page_no]['layout_dets'],
liukaiwen's avatar
liukaiwen committed
249
250
251
252
253
254
255
256
                    ),
                )
            )
        )

        objects = self.__reduct_overlap(
            list(
                map(
257
                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
liukaiwen's avatar
liukaiwen committed
258
                    filter(
259
260
                        lambda x: x['category_id'] == object_category_id,
                        self.__model_list[page_no]['layout_dets'],
liukaiwen's avatar
liukaiwen committed
261
262
263
264
265
266
                    ),
                )
            )
        )
        subject_object_relation_map = {}

许瑞's avatar
许瑞 committed
267
        subjects.sort(
268
            key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2
许瑞's avatar
许瑞 committed
269
        )  # get the distance !
liukaiwen's avatar
liukaiwen committed
270
271
272
273

        all_bboxes = []

        for v in subjects:
许瑞's avatar
许瑞 committed
274
275
            all_bboxes.append(
                {
276
277
278
                    'category_id': subject_category_id,
                    'bbox': v['bbox'],
                    'score': v['score'],
许瑞's avatar
许瑞 committed
279
280
                }
            )
liukaiwen's avatar
liukaiwen committed
281
282

        for v in objects:
许瑞's avatar
许瑞 committed
283
284
            all_bboxes.append(
                {
285
286
287
                    'category_id': object_category_id,
                    'bbox': v['bbox'],
                    'score': v['score'],
许瑞's avatar
许瑞 committed
288
289
                }
            )
liukaiwen's avatar
liukaiwen committed
290
291
292
293
294
295
296

        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 (
297
298
                    all_bboxes[i]['category_id'] == subject_category_id
                    and all_bboxes[j]['category_id'] == subject_category_id
liukaiwen's avatar
liukaiwen committed
299
300
301
                ):
                    continue

302
                dis[i][j] = bbox_distance(all_bboxes[i]['bbox'], all_bboxes[j]['bbox'])
liukaiwen's avatar
liukaiwen committed
303
304
305
306
307
                dis[j][i] = dis[i][j]

        used = set()
        for i in range(N):
            # 求第 i 个 subject 所关联的 object
308
            if all_bboxes[i]['category_id'] != subject_category_id:
liukaiwen's avatar
liukaiwen committed
309
310
311
312
313
314
315
316
317
318
319
                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(
320
                                all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
liukaiwen's avatar
liukaiwen committed
321
322
323
324
325
326
327
                            ),
                        )
                    )
                )
                if pos_flag_count > 1:
                    continue
                if (
328
                    all_bboxes[j]['category_id'] != object_category_id
329
330
                    or j in used
                    or dis[i][j] == MAX_DIS_OF_POINT
liukaiwen's avatar
liukaiwen committed
331
332
                ):
                    continue
blue's avatar
blue committed
333
                left, right, _, _ = bbox_relative_pos(
334
                    all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
blue's avatar
blue committed
335
                )  # 由  pos_flag_count 相关逻辑保证本段逻辑准确性
许瑞's avatar
许瑞 committed
336
                if left or right:
337
                    one_way_dis = all_bboxes[i]['bbox'][2] - all_bboxes[i]['bbox'][0]
许瑞's avatar
许瑞 committed
338
                else:
339
                    one_way_dis = all_bboxes[i]['bbox'][3] - all_bboxes[i]['bbox'][1]
许瑞's avatar
许瑞 committed
340
341
                if dis[i][j] > one_way_dis:
                    continue
liukaiwen's avatar
liukaiwen committed
342
343
344
345
                arr.append((dis[i][j], j))

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

许瑞's avatar
许瑞 committed
352
353
                    candidates.append(arr[0][1])
                    seen.add(arr[0][1])
liukaiwen's avatar
liukaiwen committed
354
355
356
357
358
359
360
361
362
363

            # 已经获取初始种子
            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(
364
                                    all_bboxes[j]['bbox'], all_bboxes[k]['bbox']
liukaiwen's avatar
liukaiwen committed
365
366
367
368
369
370
371
372
373
                                ),
                            )
                        )
                    )

                    if pos_flag_count > 1:
                        continue

                    if (
374
                        all_bboxes[k]['category_id'] != object_category_id
375
376
377
                        or k in used
                        or k in seen
                        or dis[j][k] == MAX_DIS_OF_POINT
许瑞's avatar
许瑞 committed
378
                        or dis[j][k] > dis[i][j]
liukaiwen's avatar
liukaiwen committed
379
380
                    ):
                        continue
许瑞's avatar
许瑞 committed
381

liukaiwen's avatar
liukaiwen committed
382
                    is_nearest = True
383
384
                    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
385
386
                            continue

387
                        if not float_gt(dis[ni][k], dis[j][k]):
liukaiwen's avatar
liukaiwen committed
388
389
390
391
                            is_nearest = False
                            break

                    if is_nearest:
blue's avatar
blue committed
392
                        nx0, ny0, nx1, ny1 = expand_bbbox(list(seen) + [k])
393
394
395
                        n_dis = bbox_distance(
                            all_bboxes[i]['bbox'], [nx0, ny0, nx1, ny1]
                        )
blue's avatar
blue committed
396
397
                        if float_gt(dis[i][j], n_dis):
                            continue
liukaiwen's avatar
liukaiwen committed
398
399
400
401
402
403
404
405
406
                        tmp.append(k)
                        seen.add(k)

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

            # 已经获取到某个 figure 下所有的最靠近的 captions,以及最靠近这些 captions 的 captions 。
            # 先扩一下 bbox,
blue's avatar
blue committed
407
            ox0, oy0, ox1, oy1 = expand_bbbox(list(seen) + [i])
408
            ix0, iy0, ix1, iy1 = all_bboxes[i]['bbox']
liukaiwen's avatar
liukaiwen committed
409
410
411
412
413
414
415
416
417
418
419
420
421

            # 分成了 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
422
423
                    if (
                        calculate_overlap_area_in_bbox1_area_ratio(
424
                            all_bboxes[idx]['bbox'], bbox
许瑞's avatar
许瑞 committed
425
426
427
                        )
                        > CAPATION_OVERLAP_AREA_RATIO
                    ):
liukaiwen's avatar
liukaiwen committed
428
429
430
                        embed_arr.append(idx)

                if len(embed_arr) > 0:
431
432
433
434
                    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
435
436
437
438
439
440
441
442
443
444
445
446
                    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
447
448
                    if (
                        calculate_overlap_area_in_bbox1_area_ratio(
449
                            all_bboxes[j]['bbox'], caption_bbox
许瑞's avatar
许瑞 committed
450
451
452
                        )
                        > CAPATION_OVERLAP_AREA_RATIO
                    ):
liukaiwen's avatar
liukaiwen committed
453
454
455
456
457
                        used.add(j)
                        subject_object_relation_map[i].append(j)

        for i in sorted(subject_object_relation_map.keys()):
            result = {
458
459
460
                'subject_body': all_bboxes[i]['bbox'],
                'all': all_bboxes[i]['bbox'],
                'score': all_bboxes[i]['score'],
liukaiwen's avatar
liukaiwen committed
461
462
463
464
            }

            if len(subject_object_relation_map[i]) > 0:
                x0 = min(
465
                    [all_bboxes[j]['bbox'][0] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
466
467
                )
                y0 = min(
468
                    [all_bboxes[j]['bbox'][1] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
469
470
                )
                x1 = max(
471
                    [all_bboxes[j]['bbox'][2] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
472
473
                )
                y1 = max(
474
                    [all_bboxes[j]['bbox'][3] for j in subject_object_relation_map[i]]
liukaiwen's avatar
liukaiwen committed
475
                )
476
477
478
479
480
481
                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
482
483
484
485
486
487
488
489
                ]
            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]:
                total_subject_object_dis += bbox_distance(
490
                    all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
liukaiwen's avatar
liukaiwen committed
491
492
493
494
495
496
497
498
499
500
501
                )

        # 计算未匹配的 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):
502
            if all_bboxes[i]['category_id'] != object_category_id or i in used:
liukaiwen's avatar
liukaiwen committed
503
504
505
506
                continue
            candidates = []
            for j in range(N):
                if (
507
                    all_bboxes[j]['category_id'] != subject_category_id
508
                    or j in with_caption_subject
liukaiwen's avatar
liukaiwen committed
509
510
511
512
513
514
515
516
517
                ):
                    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

blue's avatar
blue committed
518
    def get_imgs(self, page_no: int):
519
520
521
        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
522
        )
523
524
525
526
527
528
529
530
531
        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
532
            }
533
534
535
536
537
538
539
540

            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
541
542

    def get_tables(
543
        self, page_no: int
liukaiwen's avatar
liukaiwen committed
544
545
546
547
548
549
550
551
    ) -> 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 = {
552
553
554
555
                '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
556
557
            }

558
559
560
561
562
            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
563
564
565
566
            ret.append(record)
        return ret

    def get_equations(self, page_no: int) -> list:  # 有坐标,也有字
567
        inline_equations = self.__get_blocks_by_type(
568
            ModelBlockTypeEnum.EMBEDDING.value, page_no, ['latex']
569
570
        )
        interline_equations = self.__get_blocks_by_type(
571
            ModelBlockTypeEnum.ISOLATED.value, page_no, ['latex']
572
573
574
575
        )
        interline_equations_blocks = self.__get_blocks_by_type(
            ModelBlockTypeEnum.ISOLATE_FORMULA.value, page_no
        )
liukaiwen's avatar
liukaiwen committed
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
        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]
593
        layout_dets = model_page_info['layout_dets']
liukaiwen's avatar
liukaiwen committed
594
        for layout_det in layout_dets:
595
            if layout_det['category_id'] == '15':
liukaiwen's avatar
liukaiwen committed
596
                span = {
597
598
                    'bbox': layout_det['bbox'],
                    'content': layout_det['text'],
liukaiwen's avatar
liukaiwen committed
599
600
601
602
603
                }
                text_spans.append(span)
        return text_spans

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

605
606
607
608
609
610
        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
611

liukaiwen's avatar
liukaiwen committed
612
613
        all_spans = []
        model_page_info = self.__model_list[page_no]
614
        layout_dets = model_page_info['layout_dets']
liukaiwen's avatar
liukaiwen committed
615
616
617
        allow_category_id_list = [3, 5, 13, 14, 15]
        """当成span拼接的"""
        #  3: 'image', # 图片
618
        #  5: 'table',       # 表格
liukaiwen's avatar
liukaiwen committed
619
620
621
622
        #  13: 'inline_equation',     # 行内公式
        #  14: 'interline_equation',      # 行间公式
        #  15: 'text',      # ocr识别文本
        for layout_det in layout_dets:
623
            category_id = layout_det['category_id']
liukaiwen's avatar
liukaiwen committed
624
            if category_id in allow_category_id_list:
625
                span = {'bbox': layout_det['bbox'], 'score': layout_det['score']}
liukaiwen's avatar
liukaiwen committed
626
                if category_id == 3:
627
                    span['type'] = ContentType.Image
liukaiwen's avatar
liukaiwen committed
628
                elif category_id == 5:
629
630
                    # 获取table模型结果
                    latex = layout_det.get("latex", None)
631
                    html = layout_det.get("html", None)
632
633
                    if latex:
                        span["latex"] = latex
634
635
                    elif html:
                        span["html"] = html
liukaiwen's avatar
liukaiwen committed
636
637
                    span["type"] = ContentType.Table
                elif category_id == 13:
638
639
                    span['content'] = layout_det['latex']
                    span['type'] = ContentType.InlineEquation
liukaiwen's avatar
liukaiwen committed
640
                elif category_id == 14:
641
642
                    span['content'] = layout_det['latex']
                    span['type'] = ContentType.InterlineEquation
liukaiwen's avatar
liukaiwen committed
643
                elif category_id == 15:
644
645
                    span['content'] = layout_det['text']
                    span['type'] = ContentType.Text
liukaiwen's avatar
liukaiwen committed
646
                all_spans.append(span)
647
        return remove_duplicate_spans(all_spans)
liukaiwen's avatar
liukaiwen committed
648
649
650
651
652
653
654
655
656

    def get_page_size(self, page_no: int):  # 获取页面宽高
        # 获取当前页的page对象
        page = self.__docs[page_no]
        # 获取当前页的宽高
        page_w = page.rect.width
        page_h = page.rect.height
        return page_w, page_h

657
658
659
    def __get_blocks_by_type(
        self, type: int, page_no: int, extra_col: list[str] = []
    ) -> list:
liukaiwen's avatar
liukaiwen committed
660
661
        blocks = []
        for page_dict in self.__model_list:
662
663
664
            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
665
666
667
            if page_no != page_number:
                continue
            for item in layout_dets:
668
669
                category_id = item.get('category_id', -1)
                bbox = item.get('bbox', None)
liukaiwen's avatar
liukaiwen committed
670

liukaiwen's avatar
liukaiwen committed
671
                if category_id == type:
672
                    block = {
673
674
                        'bbox': bbox,
                        'score': item.get('score'),
675
                    }
liukaiwen's avatar
liukaiwen committed
676
677
678
679
680
                    for col in extra_col:
                        block[col] = item.get(col, None)
                    blocks.append(block)
        return blocks

许瑞's avatar
许瑞 committed
681
682
683
    def get_model_list(self, page_no):
        return self.__model_list[page_no]

684

685
686
if __name__ == '__main__':
    drw = DiskReaderWriter(r'D:/project/20231108code-clean')
liukaiwen's avatar
liukaiwen committed
687
    if 0:
688
689
        pdf_file_path = r'linshixuqiu\19983-00.pdf'
        model_file_path = r'linshixuqiu\19983-00_new.json'
liukaiwen's avatar
liukaiwen committed
690
691
692
        pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
        model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
        model_list = json.loads(model_json_txt)
693
694
        write_path = r'D:\project\20231108code-clean\linshixuqiu\19983-00'
        img_bucket_path = 'imgs'
liukaiwen's avatar
liukaiwen committed
695
        img_writer = DiskReaderWriter(join_path(write_path, img_bucket_path))
696
        pdf_docs = fitz.open('pdf', pdf_bytes)
liukaiwen's avatar
liukaiwen committed
697
698
699
700
        magic_model = MagicModel(model_list, pdf_docs)

    if 1:
        model_list = json.loads(
701
            drw.read('/opt/data/pdf/20240418/j.chroma.2009.03.042.json')
liukaiwen's avatar
liukaiwen committed
702
703
        )
        pdf_bytes = drw.read(
704
            '/opt/data/pdf/20240418/j.chroma.2009.03.042.pdf', AbsReaderWriter.MODE_BIN
liukaiwen's avatar
liukaiwen committed
705
        )
706
        pdf_docs = fitz.open('pdf', pdf_bytes)
liukaiwen's avatar
liukaiwen committed
707
708
709
        magic_model = MagicModel(model_list, pdf_docs)
        for i in range(7):
            print(magic_model.get_imgs(i))