draw_bbox.py 1.35 KB
Newer Older
1
from magic_pdf.libs.commons import fitz  # PyMuPDF
liukaiwen's avatar
lkw  
liukaiwen committed
2
3
4
5
6
7
8
9
10
11
12
13
from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
import json





def read_json_file(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    return data

赵小蒙's avatar
赵小蒙 committed
14
15

# PDF文件路径
liukaiwen's avatar
lkw  
liukaiwen committed
16
pdf_path = "D:\\projects\\Magic-PDF\\ocr_demo\\ocr_0_org.pdf"
赵小蒙's avatar
赵小蒙 committed
17
18
19
20

doc = fitz.open(pdf_path)  # Open the PDF
# 你的数据
data = [[[-2, 0, 603, 80, 24]], [[-3, 0, 602, 80, 24]]]
liukaiwen's avatar
lkw  
liukaiwen committed
21
22
23
24
25
26
27
28
29
30
31
32
33
34
ocr_json_file_path = r"D:\projects\Magic-PDF\ocr_demo\ocr_0.json"
ocr_pdf_info = read_json_file(ocr_json_file_path)
pdf_info_dict = parse_pdf_by_ocr(ocr_pdf_info)
data_list = []
for page in pdf_info_dict.values():
    page_list = []
    blocks = page.get("preproc_blocks")
    for block in blocks:
        lines = block.get("lines")
        for line in lines:
            spans = line.get("spans")
            for span in spans:
                page_list.append(span["bbox"])
    data_list.append(page_list)
赵小蒙's avatar
赵小蒙 committed
35
36
37
# 对每个页面进行处理
for i, page in enumerate(doc):
    # 获取当前页面的数据
liukaiwen's avatar
lkw  
liukaiwen committed
38
    page_data = data_list[i]
赵小蒙's avatar
赵小蒙 committed
39
    for img in page_data:
liukaiwen's avatar
lkw  
liukaiwen committed
40
        x0, y0, x1, y1 = img
赵小蒙's avatar
赵小蒙 committed
41
42
43
44
        rect_coords = fitz.Rect(x0, y0, x1, y1)  # Define the rectangle
        page.draw_rect(rect_coords, color=(1, 0, 0), fill=None, width=1.5, overlay=True)  # Draw the rectangle

# Save the PDF
liukaiwen's avatar
lkw  
liukaiwen committed
45
doc.save("D:\\projects\\Magic-PDF\\ocr_demo\\ocr_0_new1.pdf")