mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-11-02 21:42:09 +08:00
328 lines
9.9 KiB
Python
328 lines
9.9 KiB
Python
from magic_pdf.libs.commons import fitz
|
|
import os
|
|
from magic_pdf.libs.coordinate_transform import get_scale_ratio
|
|
|
|
|
|
def draw_model_output(
|
|
raw_pdf_doc: fitz.Document, paras_dict_arr: list[dict], save_path: str
|
|
):
|
|
"""
|
|
在page上画出bbox，保存到save_path
|
|
"""
|
|
"""
|
|
|
|
# {0: 'title', # 标题
|
|
# 1: 'figure', # 图片
|
|
# 2: 'plain text', # 文本
|
|
# 3: 'header', # 页眉
|
|
# 4: 'page number', # 页码
|
|
# 5: 'footnote', # 脚注
|
|
# 6: 'footer', # 页脚
|
|
# 7: 'table', # 表格
|
|
# 8: 'table caption', # 表格描述
|
|
# 9: 'figure caption', # 图片描述
|
|
# 10: 'equation', # 公式
|
|
# 11: 'full column', # 单栏
|
|
# 12: 'sub column', # 多栏
|
|
# 13: 'embedding', # 嵌入公式
|
|
# 14: 'isolated'} # 单行公式
|
|
|
|
"""
|
|
|
|
color_map = {
|
|
"body": fitz.pdfcolor["green"],
|
|
"non_body": fitz.pdfcolor["red"],
|
|
}
|
|
"""
|
|
{"layout_dets": [], "subfield_dets": [], "page_info": {"page_no": 22, "height": 1650, "width": 1275}}
|
|
"""
|
|
for i, page in enumerate(raw_pdf_doc):
|
|
v = paras_dict_arr[i]
|
|
page_idx = v["page_info"]["page_no"]
|
|
width = v["page_info"]["width"]
|
|
height = v["page_info"]["height"]
|
|
|
|
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
|
|
paras_dict_arr[i], page
|
|
)
|
|
|
|
for order, block in enumerate(v["layout_dets"]):
|
|
L = block["poly"][0] / horizontal_scale_ratio
|
|
U = block["poly"][1] / vertical_scale_ratio
|
|
R = block["poly"][2] / horizontal_scale_ratio
|
|
D = block["poly"][5] / vertical_scale_ratio
|
|
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
|
|
# R += pageL
|
|
# U += pageU
|
|
# D += pageU
|
|
L, R = min(L, R), max(L, R)
|
|
U, D = min(U, D), max(U, D)
|
|
bbox = [L, U, R, D]
|
|
color = color_map["body"]
|
|
if block["category_id"] in (3, 4, 5, 6, 0):
|
|
color = color_map["non_body"]
|
|
|
|
rect = fitz.Rect(bbox)
|
|
page.draw_rect(rect, fill=None, width=0.5, overlay=True, color=color)
|
|
|
|
parent_dir = os.path.dirname(save_path)
|
|
if not os.path.exists(parent_dir):
|
|
os.makedirs(parent_dir)
|
|
raw_pdf_doc.save(save_path)
|
|
|
|
|
|
def debug_show_bbox(
|
|
raw_pdf_doc: fitz.Document,
|
|
page_idx: int,
|
|
bboxes: list,
|
|
droped_bboxes: list,
|
|
expect_drop_bboxes: list,
|
|
save_path: str,
|
|
expected_page_id: int,
|
|
):
|
|
"""
|
|
以覆盖的方式写个临时的pdf，用于debug
|
|
"""
|
|
if page_idx != expected_page_id:
|
|
return
|
|
|
|
if os.path.exists(save_path):
|
|
# 删除已经存在的文件
|
|
os.remove(save_path)
|
|
# 创建一个新的空白 PDF 文件
|
|
doc = fitz.open("")
|
|
|
|
width = raw_pdf_doc[page_idx].rect.width
|
|
height = raw_pdf_doc[page_idx].rect.height
|
|
new_page = doc.new_page(width=width, height=height)
|
|
|
|
shape = new_page.new_shape()
|
|
for bbox in bboxes:
|
|
# 原始box画上去
|
|
rect = fitz.Rect(*bbox[0:4])
|
|
shape = new_page.new_shape()
|
|
shape.draw_rect(rect)
|
|
shape.finish(
|
|
color=fitz.pdfcolor["red"], fill=fitz.pdfcolor["blue"], fill_opacity=0.2
|
|
)
|
|
shape.finish()
|
|
shape.commit()
|
|
|
|
for bbox in droped_bboxes:
|
|
# 原始box画上去
|
|
rect = fitz.Rect(*bbox[0:4])
|
|
shape = new_page.new_shape()
|
|
shape.draw_rect(rect)
|
|
shape.finish(color=None, fill=fitz.pdfcolor["yellow"], fill_opacity=0.2)
|
|
shape.finish()
|
|
shape.commit()
|
|
|
|
for bbox in expect_drop_bboxes:
|
|
# 原始box画上去
|
|
rect = fitz.Rect(*bbox[0:4])
|
|
shape = new_page.new_shape()
|
|
shape.draw_rect(rect)
|
|
shape.finish(color=fitz.pdfcolor["red"], fill=None)
|
|
shape.finish()
|
|
shape.commit()
|
|
|
|
# shape.insert_textbox(fitz.Rect(200, 0, 600, 20), f"total bboxes: {len(bboxes)}", fontname="helv", fontsize=12,
|
|
# color=(0, 0, 0))
|
|
# shape.finish(color=fitz.pdfcolor['black'])
|
|
# shape.commit()
|
|
|
|
parent_dir = os.path.dirname(save_path)
|
|
if not os.path.exists(parent_dir):
|
|
os.makedirs(parent_dir)
|
|
|
|
doc.save(save_path)
|
|
doc.close()
|
|
|
|
|
|
def debug_show_page(
|
|
page,
|
|
bboxes1: list,
|
|
bboxes2: list,
|
|
bboxes3: list,
|
|
):
|
|
save_path = "./tmp/debug.pdf"
|
|
if os.path.exists(save_path):
|
|
# 删除已经存在的文件
|
|
os.remove(save_path)
|
|
# 创建一个新的空白 PDF 文件
|
|
doc = fitz.open("")
|
|
|
|
width = page.rect.width
|
|
height = page.rect.height
|
|
new_page = doc.new_page(width=width, height=height)
|
|
|
|
shape = new_page.new_shape()
|
|
for bbox in bboxes1:
|
|
# 原始box画上去
|
|
rect = fitz.Rect(*bbox[0:4])
|
|
shape = new_page.new_shape()
|
|
shape.draw_rect(rect)
|
|
shape.finish(
|
|
color=fitz.pdfcolor["red"], fill=fitz.pdfcolor["blue"], fill_opacity=0.2
|
|
)
|
|
shape.finish()
|
|
shape.commit()
|
|
|
|
for bbox in bboxes2:
|
|
# 原始box画上去
|
|
rect = fitz.Rect(*bbox[0:4])
|
|
shape = new_page.new_shape()
|
|
shape.draw_rect(rect)
|
|
shape.finish(color=None, fill=fitz.pdfcolor["yellow"], fill_opacity=0.2)
|
|
shape.finish()
|
|
shape.commit()
|
|
|
|
for bbox in bboxes3:
|
|
# 原始box画上去
|
|
rect = fitz.Rect(*bbox[0:4])
|
|
shape = new_page.new_shape()
|
|
shape.draw_rect(rect)
|
|
shape.finish(color=fitz.pdfcolor["red"], fill=None)
|
|
shape.finish()
|
|
shape.commit()
|
|
|
|
parent_dir = os.path.dirname(save_path)
|
|
if not os.path.exists(parent_dir):
|
|
os.makedirs(parent_dir)
|
|
|
|
doc.save(save_path)
|
|
doc.close()
|
|
|
|
|
|
def draw_layout_bbox_on_page(
|
|
raw_pdf_doc: fitz.Document, paras_dict: dict, header, footer, pdf_path: str
|
|
):
|
|
"""
|
|
在page上画出bbox，保存到save_path
|
|
"""
|
|
# 检查文件是否存在
|
|
is_new_pdf = False
|
|
if os.path.exists(pdf_path):
|
|
# 打开现有的 PDF 文件
|
|
doc = fitz.open(pdf_path)
|
|
else:
|
|
# 创建一个新的空白 PDF 文件
|
|
is_new_pdf = True
|
|
doc = fitz.open("")
|
|
|
|
for k, v in paras_dict.items():
|
|
page_idx = v["page_idx"]
|
|
layouts = v["layout_bboxes"]
|
|
page = doc[page_idx]
|
|
shape = page.new_shape()
|
|
for order, layout in enumerate(layouts):
|
|
border_offset = 1
|
|
rect_box = layout["layout_bbox"]
|
|
layout_label = layout["layout_label"]
|
|
fill_color = fitz.pdfcolor["pink"] if layout_label == "U" else None
|
|
rect_box = [
|
|
rect_box[0] + 1,
|
|
rect_box[1] - border_offset,
|
|
rect_box[2] - 1,
|
|
rect_box[3] + border_offset,
|
|
]
|
|
rect = fitz.Rect(*rect_box)
|
|
shape.draw_rect(rect)
|
|
shape.finish(color=fitz.pdfcolor["red"], fill=fill_color, fill_opacity=0.4)
|
|
"""
|
|
draw order text on layout box
|
|
"""
|
|
font_size = 10
|
|
shape.insert_text(
|
|
(rect_box[0] + 1, rect_box[1] + font_size),
|
|
f"{order}",
|
|
fontsize=font_size,
|
|
color=(0, 0, 0),
|
|
)
|
|
|
|
"""画上footer header"""
|
|
if header:
|
|
shape.draw_rect(fitz.Rect(header))
|
|
shape.finish(color=None, fill=fitz.pdfcolor["black"], fill_opacity=0.2)
|
|
if footer:
|
|
shape.draw_rect(fitz.Rect(footer))
|
|
shape.finish(color=None, fill=fitz.pdfcolor["black"], fill_opacity=0.2)
|
|
|
|
shape.commit()
|
|
|
|
if is_new_pdf:
|
|
doc.save(pdf_path)
|
|
else:
|
|
doc.saveIncr()
|
|
doc.close()
|
|
|
|
|
|
@DeprecationWarning
|
|
def draw_layout_on_page(
|
|
raw_pdf_doc: fitz.Document, page_idx: int, page_layout: list, pdf_path: str
|
|
):
|
|
"""
|
|
把layout的box用红色边框花在pdf_path的page_idx上
|
|
"""
|
|
|
|
def draw(shape, layout, fill_color=fitz.pdfcolor["pink"]):
|
|
border_offset = 1
|
|
rect_box = layout["layout_bbox"]
|
|
layout_label = layout["layout_label"]
|
|
sub_layout = layout["sub_layout"]
|
|
if len(sub_layout) == 0:
|
|
fill_color = fill_color if layout_label == "U" else None
|
|
rect_box = [
|
|
rect_box[0] + 1,
|
|
rect_box[1] - border_offset,
|
|
rect_box[2] - 1,
|
|
rect_box[3] + border_offset,
|
|
]
|
|
rect = fitz.Rect(*rect_box)
|
|
shape.draw_rect(rect)
|
|
shape.finish(color=fitz.pdfcolor["red"], fill=fill_color, fill_opacity=0.2)
|
|
# if layout_label=='U':
|
|
# bad_boxes = layout.get("bad_boxes", [])
|
|
# for bad_box in bad_boxes:
|
|
# rect = fitz.Rect(*bad_box)
|
|
# shape.draw_rect(rect)
|
|
# shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['red'], fill_opacity=0.2)
|
|
# else:
|
|
# rect = fitz.Rect(*rect_box)
|
|
# shape.draw_rect(rect)
|
|
# shape.finish(color=fitz.pdfcolor['blue'])
|
|
|
|
for sub_layout in sub_layout:
|
|
draw(shape, sub_layout)
|
|
shape.commit()
|
|
|
|
# 检查文件是否存在
|
|
is_new_pdf = False
|
|
if os.path.exists(pdf_path):
|
|
# 打开现有的 PDF 文件
|
|
doc = fitz.open(pdf_path)
|
|
else:
|
|
# 创建一个新的空白 PDF 文件
|
|
is_new_pdf = True
|
|
doc = fitz.open("")
|
|
|
|
page = doc[page_idx]
|
|
shape = page.new_shape()
|
|
for order, layout in enumerate(page_layout):
|
|
draw(shape, layout, fitz.pdfcolor["yellow"])
|
|
|
|
# shape.insert_textbox(fitz.Rect(200, 0, 600, 20), f"total bboxes: {len(layout)}", fontname="helv", fontsize=12,
|
|
# color=(0, 0, 0))
|
|
# shape.finish(color=fitz.pdfcolor['black'])
|
|
# shape.commit()
|
|
|
|
parent_dir = os.path.dirname(pdf_path)
|
|
if not os.path.exists(parent_dir):
|
|
os.makedirs(parent_dir)
|
|
|
|
if is_new_pdf:
|
|
doc.save(pdf_path)
|
|
else:
|
|
doc.saveIncr()
|
|
doc.close()
|