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

赵小蒙'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
15
16
17
        if fill_config:
            page.draw_rect(rect_coords, color=None, fill=new_rgb, fill_opacity=0.3, width=0.5, overlay=True)  # Draw the rectangle
        else:
            page.draw_rect(rect_coords, color=new_rgb, fill=None, fill_opacity=1, width=0.5, overlay=True)  # Draw the rectangle
赵小蒙's avatar
赵小蒙 committed
18

19

赵小蒙's avatar
赵小蒙 committed
20
def draw_bbox_with_number(i, bbox_list, page, rgb_config, fill_config):
赵小蒙's avatar
赵小蒙 committed
21
22
23
24
25
26
27
28
    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
29
30
31
32
33
        if fill_config:
            page.draw_rect(rect_coords, color=None, fill=new_rgb, fill_opacity=0.3, width=0.5, overlay=True)  # Draw the rectangle
        else:
            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
34
35


赵小蒙's avatar
赵小蒙 committed
36
def draw_layout_bbox(pdf_info, pdf_bytes, out_path):
37
    layout_bbox_list = []
赵小蒙's avatar
赵小蒙 committed
38
    blocks_bbox_list = []
赵小蒙's avatar
赵小蒙 committed
39
    dropped_bbox_list = []
赵小蒙's avatar
赵小蒙 committed
40
    for page in pdf_info:
赵小蒙's avatar
赵小蒙 committed
41
42
        page_layout_list = []
        page_dropped_list = []
赵小蒙's avatar
赵小蒙 committed
43
        page_blocks_bbox_list = []
44
        for layout in page['layout_bboxes']:
赵小蒙's avatar
赵小蒙 committed
45
46
            page_layout_list.append(layout['layout_bbox'])
        layout_bbox_list.append(page_layout_list)
赵小蒙's avatar
赵小蒙 committed
47
48
        for dropped_bbox in page['discarded_blocks']:
            page_dropped_list.append(dropped_bbox['bbox'])
赵小蒙's avatar
赵小蒙 committed
49
        dropped_bbox_list.append(page_dropped_list)
赵小蒙's avatar
赵小蒙 committed
50
51
52
        for block in page['para_blocks']:
            page_blocks_bbox_list.append(block['bbox'])
        blocks_bbox_list.append(page_blocks_bbox_list)
53
54
    pdf_docs = fitz.open("pdf", pdf_bytes)
    for i, page in enumerate(pdf_docs):
赵小蒙's avatar
赵小蒙 committed
55
56
57
        draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0], False)
        draw_bbox_without_number(i, dropped_bbox_list, page, [0, 255, 0], True)
        draw_bbox_without_number(i, blocks_bbox_list, page, [0, 0, 255], True)
58
    # Save the PDF
59
    pdf_docs.save(f"{out_path}/layout.pdf")
60

61
def draw_text_bbox(pdf_info_dict, pdf_bytes, out_path):
62
63
    text_list = []
    inline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
64
    interline_equation_list = []
65
66
67
    for page in pdf_info_dict.values():
        page_text_list = []
        page_inline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
68
        page_interline_equation_list = []
赵小蒙's avatar
赵小蒙 committed
69
        for block in page['para_blocks']:
70
71
            for line in block['lines']:
                for span in line['spans']:
赵小蒙's avatar
赵小蒙 committed
72
                    if span['type'] == ContentType.Text:
73
                        page_text_list.append(span['bbox'])
赵小蒙's avatar
赵小蒙 committed
74
                    elif span['type'] == ContentType.InlineEquation:
75
                        page_inline_equation_list.append(span['bbox'])
赵小蒙's avatar
赵小蒙 committed
76
77
                    elif span['type'] == ContentType.InterlineEquation:
                        page_interline_equation_list.append(span['bbox'])
78
79
        text_list.append(page_text_list)
        inline_equation_list.append(page_inline_equation_list)
赵小蒙's avatar
赵小蒙 committed
80
        interline_equation_list.append(page_interline_equation_list)
81
82
    pdf_docs = fitz.open("pdf", pdf_bytes)
    for i, page in enumerate(pdf_docs):
83
        # 获取当前页面的数据
赵小蒙's avatar
赵小蒙 committed
84
85
        draw_bbox_without_number(i, text_list, page, [255, 0, 0])
        draw_bbox_without_number(i, inline_equation_list, page, [0, 255, 0])
赵小蒙's avatar
赵小蒙 committed
86
        draw_bbox_without_number(i, interline_equation_list, page, [0, 0, 255])
87
88

    # Save the PDF
89
    pdf_docs.save(f"{out_path}/text.pdf")