draw_bbox.py 9.23 KB
Newer Older
1
from magic_pdf.libs.commons import fitz  # PyMuPDF
赵小蒙's avatar
赵小蒙 committed
2
from magic_pdf.libs.ocr_content_type import ContentType, BlockType
赵小蒙's avatar
赵小蒙 committed
3

赵小蒙's avatar
赵小蒙 committed
4

赵小蒙's avatar
赵小蒙 committed
5
def draw_bbox_without_number(i, bbox_list, page, rgb_config, fill_config):
6
7
8
9
10
11
12
    new_rgb = []
    for item in rgb_config:
        item = float(item) / 255
        new_rgb.append(item)
    page_data = bbox_list[i]
    for bbox in page_data:
        x0, y0, x1, y1 = bbox
赵小蒙's avatar
赵小蒙 committed
13
        rect_coords = fitz.Rect(x0, y0, x1, y1)  # Define the rectangle
赵小蒙's avatar
赵小蒙 committed
14
        if fill_config:
许瑞's avatar
许瑞 committed
15
16
17
18
19
20
21
22
            page.draw_rect(
                rect_coords,
                color=None,
                fill=new_rgb,
                fill_opacity=0.3,
                width=0.5,
                overlay=True,
            )  # Draw the rectangle
赵小蒙's avatar
赵小蒙 committed
23
        else:
许瑞's avatar
许瑞 committed
24
25
26
27
28
29
30
31
            page.draw_rect(
                rect_coords,
                color=new_rgb,
                fill=None,
                fill_opacity=1,
                width=0.5,
                overlay=True,
            )  # Draw the rectangle
赵小蒙's avatar
赵小蒙 committed
32

33

赵小蒙's avatar
赵小蒙 committed
34
def draw_bbox_with_number(i, bbox_list, page, rgb_config, fill_config):
赵小蒙's avatar
赵小蒙 committed
35
36
37
38
39
40
41
42
    new_rgb = []
    for item in rgb_config:
        item = float(item) / 255
        new_rgb.append(item)
    page_data = bbox_list[i]
    for j, bbox in enumerate(page_data):
        x0, y0, x1, y1 = bbox
        rect_coords = fitz.Rect(x0, y0, x1, y1)  # Define the rectangle
赵小蒙's avatar
赵小蒙 committed
43
        if fill_config:
许瑞's avatar
许瑞 committed
44
45
46
47
48
49
50
51
            page.draw_rect(
                rect_coords,
                color=None,
                fill=new_rgb,
                fill_opacity=0.3,
                width=0.5,
                overlay=True,
            )  # Draw the rectangle
赵小蒙's avatar
赵小蒙 committed
52
        else:
许瑞's avatar
许瑞 committed
53
54
55
56
57
58
59
60
61
62
63
            page.draw_rect(
                rect_coords,
                color=new_rgb,
                fill=None,
                fill_opacity=1,
                width=0.5,
                overlay=True,
            )  # Draw the rectangle
        page.insert_text(
            (x0, y0 + 10), str(j + 1), fontsize=10, color=new_rgb
        )  # Insert the index at the top left corner of the rectangle
赵小蒙's avatar
赵小蒙 committed
64
65


赵小蒙's avatar
赵小蒙 committed
66
def draw_layout_bbox(pdf_info, pdf_bytes, out_path):
67
    layout_bbox_list = []
赵小蒙's avatar
赵小蒙 committed
68
    dropped_bbox_list = []
赵小蒙's avatar
赵小蒙 committed
69
    tables_list, tables_body_list, tables_caption_list, tables_footnote_list = [], [], [], []
许瑞's avatar
许瑞 committed
70
71
72
73
    imgs_list, imgs_body_list, imgs_caption_list = [], [], []
    titles_list = []
    texts_list = []
    interequations_list = []
赵小蒙's avatar
赵小蒙 committed
74
    for page in pdf_info:
赵小蒙's avatar
赵小蒙 committed
75
76
        page_layout_list = []
        page_dropped_list = []
许瑞's avatar
许瑞 committed
77
78
79
80
81
82
83
        tables, tables_body, tables_caption, tables_footnote = [], [], [], []
        imgs, imgs_body, imgs_caption = [], [], []
        titles = []
        texts = []
        interequations = []
        for layout in page["layout_bboxes"]:
            page_layout_list.append(layout["layout_bbox"])
赵小蒙's avatar
赵小蒙 committed
84
        layout_bbox_list.append(page_layout_list)
许瑞's avatar
许瑞 committed
85
86
        for dropped_bbox in page["discarded_blocks"]:
            page_dropped_list.append(dropped_bbox["bbox"])
赵小蒙's avatar
赵小蒙 committed
87
        dropped_bbox_list.append(page_dropped_list)
许瑞's avatar
许瑞 committed
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        for block in page["para_blocks"]:
            bbox = block["bbox"]
            if block["type"] == BlockType.Table:
                tables.append(bbox)
                for nested_block in block["blocks"]:
                    bbox = nested_block["bbox"]
                    if nested_block["type"] == BlockType.TableBody:
                        tables_body.append(bbox)
                    elif nested_block["type"] == BlockType.TableCaption:
                        tables_caption.append(bbox)
                    elif nested_block["type"] == BlockType.TableFootnote:
                        tables_footnote.append(bbox)
            elif block["type"] == BlockType.Image:
                imgs.append(bbox)
                for nested_block in block["blocks"]:
                    bbox = nested_block["bbox"]
                    if nested_block["type"] == BlockType.ImageBody:
                        imgs_body.append(bbox)
                    elif nested_block["type"] == BlockType.ImageCaption:
                        imgs_caption.append(bbox)
            elif block["type"] == BlockType.Title:
                titles.append(bbox)
            elif block["type"] == BlockType.Text:
                texts.append(bbox)
            elif block["type"] == BlockType.InterlineEquation:
                interequations.append(bbox)
        tables_list.append(tables)
        tables_body_list.append(tables_body)
        tables_caption_list.append(tables_caption)
        tables_footnote_list.append(tables_footnote)
        imgs_list.append(imgs)
        imgs_body_list.append(imgs_body)
        imgs_caption_list.append(imgs_caption)
        titles_list.append(titles)
        texts_list.append(texts)
        interequations_list.append(interequations)

125
126
    pdf_docs = fitz.open("pdf", pdf_bytes)
    for i, page in enumerate(pdf_docs):
赵小蒙's avatar
赵小蒙 committed
127
        draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0], False)
赵小蒙's avatar
赵小蒙 committed
128
        draw_bbox_without_number(i, dropped_bbox_list, page, [158, 158, 158], True)
许瑞's avatar
许瑞 committed
129
130
131
132
133
134
135
136
137
        draw_bbox_without_number(i, tables_list, page, [153, 153, 0], True)  # color !
        draw_bbox_without_number(i, tables_body_list, page, [204, 204, 0], True)
        draw_bbox_without_number(i, tables_caption_list, page, [255, 255, 102], True)
        draw_bbox_without_number(i, tables_footnote_list, page, [229, 255, 204], True)
        draw_bbox_without_number(i, imgs_list, page, [51, 102, 0], True)
        draw_bbox_without_number(i, imgs_body_list, page, [153, 255, 51], True)
        draw_bbox_without_number(i, imgs_caption_list, page, [102, 178, 255], True)
        draw_bbox_without_number(i, titles_list, page, [102, 102, 255], True)
        draw_bbox_without_number(i, texts_list, page, [153, 0, 76], True)
赵小蒙's avatar
赵小蒙 committed
138
        draw_bbox_without_number(i, interequations_list, page, [0, 255, 0], True)
许瑞's avatar
许瑞 committed
139

140
    # Save the PDF
141
    pdf_docs.save(f"{out_path}/layout.pdf")
142

许瑞's avatar
许瑞 committed
143

赵小蒙's avatar
赵小蒙 committed
144
def draw_span_bbox(pdf_info, pdf_bytes, out_path):
145
146
    text_list = []
    inline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
147
    interline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
148
149
    image_list = []
    table_list = []
赵小蒙's avatar
赵小蒙 committed
150
    dropped_list = []
赵小蒙's avatar
赵小蒙 committed
151
    for page in pdf_info:
152
153
        page_text_list = []
        page_inline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
154
        page_interline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
155
156
        page_image_list = []
        page_table_list = []
赵小蒙's avatar
赵小蒙 committed
157
158
159
160
161
162
163
164
165
        page_dropped_list = []
        # 构造dropped_list
        for block in page["discarded_blocks"]:
            if block["type"] == BlockType.Discarded:
                for line in block["lines"]:
                    for span in line["spans"]:
                        page_dropped_list.append(span["bbox"])
        dropped_list.append(page_dropped_list)
        # 构造其余useful_list
许瑞's avatar
许瑞 committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
        for block in page["para_blocks"]:
            if block["type"] in [
                BlockType.Text,
                BlockType.Title,
                BlockType.InterlineEquation,
            ]:
                for line in block["lines"]:
                    for span in line["spans"]:
                        if span["type"] == ContentType.Text:
                            page_text_list.append(span["bbox"])
                        elif span["type"] == ContentType.InlineEquation:
                            page_inline_equation_list.append(span["bbox"])
                        elif span["type"] == ContentType.InterlineEquation:
                            page_interline_equation_list.append(span["bbox"])
                        elif span["type"] == ContentType.Image:
                            page_image_list.append(span["bbox"])
                        elif span["type"] == ContentType.Table:
                            page_table_list.append(span["bbox"])
            elif block["type"] in [BlockType.Image, BlockType.Table]:
赵小蒙's avatar
赵小蒙 committed
185
                for sub_block in block["blocks"]:
许瑞's avatar
许瑞 committed
186
187
188
189
190
191
192
193
194
195
196
197
                    for line in sub_block["lines"]:
                        for span in line["spans"]:
                            if span["type"] == ContentType.Text:
                                page_text_list.append(span["bbox"])
                            elif span["type"] == ContentType.InlineEquation:
                                page_inline_equation_list.append(span["bbox"])
                            elif span["type"] == ContentType.InterlineEquation:
                                page_interline_equation_list.append(span["bbox"])
                            elif span["type"] == ContentType.Image:
                                page_image_list.append(span["bbox"])
                            elif span["type"] == ContentType.Table:
                                page_table_list.append(span["bbox"])
198
199
        text_list.append(page_text_list)
        inline_equation_list.append(page_inline_equation_list)
赵小蒙's avatar
赵小蒙 committed
200
        interline_equation_list.append(page_interline_equation_list)
赵小蒙's avatar
赵小蒙 committed
201
202
        image_list.append(page_image_list)
        table_list.append(page_table_list)
203
204
    pdf_docs = fitz.open("pdf", pdf_bytes)
    for i, page in enumerate(pdf_docs):
205
        # 获取当前页面的数据
赵小蒙's avatar
赵小蒙 committed
206
207
208
209
210
        draw_bbox_without_number(i, text_list, page, [255, 0, 0], False)
        draw_bbox_without_number(i, inline_equation_list, page, [0, 255, 0], False)
        draw_bbox_without_number(i, interline_equation_list, page, [0, 0, 255], False)
        draw_bbox_without_number(i, image_list, page, [255, 204, 0], False)
        draw_bbox_without_number(i, table_list, page, [204, 0, 255], False)
赵小蒙's avatar
赵小蒙 committed
211
        draw_bbox_without_number(i, dropped_list, page, [158, 158, 158], False)
212
213

    # Save the PDF
赵小蒙's avatar
赵小蒙 committed
214
    pdf_docs.save(f"{out_path}/spans.pdf")