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()