ocr_detect_all_bboxes.py 4.13 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
2
from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio, calculate_overlap_area_in_bbox1_area_ratio, \
    calculate_iou
赵小蒙's avatar
赵小蒙 committed
3
from magic_pdf.libs.drop_tag import DropTag
赵小蒙's avatar
赵小蒙 committed
4
5
6
from magic_pdf.libs.ocr_content_type import BlockType


赵小蒙's avatar
赵小蒙 committed
7
8
9
10
11
def ocr_prepare_bboxes_for_layout_split(img_blocks, table_blocks, discarded_blocks, text_blocks,
                                        title_blocks, interline_equation_blocks, page_w, page_h):
    all_bboxes = []
    for image in img_blocks:
        x0, y0, x1, y1 = image['bbox']
赵小蒙's avatar
赵小蒙 committed
12
        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Image, None, None, None, None])
赵小蒙's avatar
赵小蒙 committed
13
14
15

    for table in table_blocks:
        x0, y0, x1, y1 = table['bbox']
赵小蒙's avatar
赵小蒙 committed
16
        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Table, None, None, None, None])
赵小蒙's avatar
赵小蒙 committed
17
18
19

    for text in text_blocks:
        x0, y0, x1, y1 = text['bbox']
赵小蒙's avatar
赵小蒙 committed
20
        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Text, None, None, None, None])
赵小蒙's avatar
赵小蒙 committed
21
22
23

    for title in title_blocks:
        x0, y0, x1, y1 = title['bbox']
赵小蒙's avatar
赵小蒙 committed
24
        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Title, None, None, None, None])
赵小蒙's avatar
赵小蒙 committed
25
26
27

    for interline_equation in interline_equation_blocks:
        x0, y0, x1, y1 = interline_equation['bbox']
赵小蒙's avatar
赵小蒙 committed
28
        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.InterlineEquation, None, None, None, None])
赵小蒙's avatar
赵小蒙 committed
29

赵小蒙's avatar
赵小蒙 committed
30
    '''block嵌套问题解决'''
31
    '''文本框与标题框重叠,优先信任文本框'''
赵小蒙's avatar
赵小蒙 committed
32
33
34
35
36
37
    all_bboxes = fix_text_overlap_title_blocks(all_bboxes)
    '''任何框体与舍弃框重叠,优先信任舍弃框'''
    all_bboxes = remove_need_drop_blocks(all_bboxes, discarded_blocks)
    '''经过以上处理后,还存在大框套小框的情况,则删除小框'''
    all_bboxes = remove_overlaps_min_blocks(all_bboxes)

赵小蒙's avatar
赵小蒙 committed
38
39
40
41
    '''discarded_blocks中只保留宽度超过1/3页面宽度的,高度超过10的,处于页面下半50%区域的(限定footnote)'''
    for discarded in discarded_blocks:
        x0, y0, x1, y1 = discarded['bbox']
        if (x1 - x0) > (page_w / 3) and (y1 - y0) > 10 and y0 > (page_h / 2):
赵小蒙's avatar
赵小蒙 committed
42
            all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Footnote, None, None, None, None])
赵小蒙's avatar
赵小蒙 committed
43

赵小蒙's avatar
赵小蒙 committed
44
45
46
47
48
49
50
51
52
53
54
55
56
    return all_bboxes


def fix_text_overlap_title_blocks(all_bboxes):
    # 先提取所有text和title block
    text_blocks = []
    for block in all_bboxes:
        if block[7] == BlockType.Text:
            text_blocks.append(block)
    title_blocks = []
    for block in all_bboxes:
        if block[7] == BlockType.Title:
            title_blocks.append(block)
赵小蒙's avatar
赵小蒙 committed
57

赵小蒙's avatar
赵小蒙 committed
58
59
    for text_block in text_blocks:
        for title_block in title_blocks:
赵小蒙's avatar
赵小蒙 committed
60
61
            text_block_bbox = text_block[:4]
            title_block_bbox = title_block[:4]
赵小蒙's avatar
赵小蒙 committed
62
            if calculate_iou(text_block_bbox, title_block_bbox) > 0.8:
63
                all_bboxes.remove(title_block)
赵小蒙's avatar
赵小蒙 committed
64
65
66
67
68

    return all_bboxes


def remove_need_drop_blocks(all_bboxes, discarded_blocks):
赵小蒙's avatar
赵小蒙 committed
69
70
    need_remove = []
    for block in all_bboxes:
赵小蒙's avatar
赵小蒙 committed
71
        for discarded_block in discarded_blocks:
赵小蒙's avatar
赵小蒙 committed
72
            block_bbox = block[:4]
赵小蒙's avatar
赵小蒙 committed
73
            if calculate_overlap_area_in_bbox1_area_ratio(block_bbox, discarded_block['bbox']) > 0.6:
赵小蒙's avatar
赵小蒙 committed
74
75
76
                if block not in need_remove:
                    need_remove.append(block)
                    break
赵小蒙's avatar
赵小蒙 committed
77

赵小蒙's avatar
赵小蒙 committed
78
79
80
    if len(need_remove) > 0:
        for block in need_remove:
            all_bboxes.remove(block)
赵小蒙's avatar
赵小蒙 committed
81
82
    return all_bboxes

赵小蒙's avatar
赵小蒙 committed
83
84
85

def remove_overlaps_min_blocks(all_bboxes):
    #  删除重叠blocks中较小的那些
赵小蒙's avatar
赵小蒙 committed
86
87
88
    need_remove = []
    for block1 in all_bboxes:
        for block2 in all_bboxes:
赵小蒙's avatar
赵小蒙 committed
89
            if block1 != block2:
赵小蒙's avatar
赵小蒙 committed
90
91
                block1_bbox = block1[:4]
                block2_bbox = block2[:4]
赵小蒙's avatar
赵小蒙 committed
92
                overlap_box = get_minbox_if_overlap_by_ratio(block1_bbox, block2_bbox, 0.8)
赵小蒙's avatar
赵小蒙 committed
93
                if overlap_box is not None:
赵小蒙's avatar
赵小蒙 committed
94
                    bbox_to_remove = next((block for block in all_bboxes if block[:4] == overlap_box), None)
赵小蒙's avatar
赵小蒙 committed
95
                    if bbox_to_remove is not None and bbox_to_remove not in need_remove:
赵小蒙's avatar
赵小蒙 committed
96
                        need_remove.append(bbox_to_remove)
赵小蒙's avatar
赵小蒙 committed
97

赵小蒙's avatar
赵小蒙 committed
98
99
100
    if len(need_remove) > 0:
        for block in need_remove:
            all_bboxes.remove(block)
赵小蒙's avatar
赵小蒙 committed
101
102

    return all_bboxes