Compare commits

...

62 Commits

Author SHA1 Message Date
赵小蒙
f06a32133c make多模态markdown时图片地址更改为fullpath 2024-03-15 14:01:23 +08:00
赵小蒙
59b0b0c3da make markdown时特殊符号转义 2024-03-14 18:49:54 +08:00
赵小蒙
8a2736a53f 截图增加s3上传逻辑,移除宽或高为0的spans 2024-03-14 17:37:53 +08:00
赵小蒙
0b35b73c64 删除高度或者宽度为0的spans 2024-03-14 17:31:21 +08:00
赵小蒙
9bd6294bfa ocr模式更新spark pipeline 2024-03-14 16:33:25 +08:00
赵小蒙
8a52ada337 data_type/bookid/data_source兼容处理 2024-03-14 15:58:34 +08:00
赵小蒙
26c2378271 ocr模式下content type 抽象 2024-03-14 15:25:47 +08:00
赵小蒙
b6f051d88f 在layout.pdf中绘制drop的bbox 2024-03-14 11:26:49 +08:00
赵小蒙
85587b257b 在dict中加入qa需要的字段 2024-03-13 19:04:24 +08:00
赵小蒙
b560c18fb4 fix import 2024-03-13 15:50:59 +08:00
myhloli
1cb79e7df0 Merge pull request #5 from myhloli/dev-in-line-bbox
Dev in line bbox
2024-03-13 15:20:08 +08:00
liukaiwen
8f2640826f Merge branch 'master' into dev-in-line-bbox
# Conflicts:
#	magic_pdf/pre_proc/ocr_span_list_modify.py
2024-03-13 14:57:22 +08:00
liukaiwen
21cfaf4c0d add modify inline equation y axis
add false displayed equation to inline equation
2024-03-13 14:55:29 +08:00
赵小蒙
6f7aa890c2 fix import 2024-03-13 14:26:45 +08:00
myhloli
f9310954bd Merge pull request #4 from myhloli/dev-in-line-bbox
Dev in line bbox
2024-03-13 14:22:57 +08:00
liukaiwen
64d67b5c58 add modify inline equation y axis
add false displayed equation to inline equation
2024-03-13 14:19:38 +08:00
liukaiwen
f5dc261da5 Merge branch 'master' into dev-in-line-bbox 2024-03-13 14:18:12 +08:00
liukaiwen
1f468bed0a add modify inline equation y axis
add false displayed equation to inline equation
2024-03-13 14:16:33 +08:00
赵小蒙
32fd7f958c 将对span的操作移动到ocr_span_list_modify,增加独占一行区块的位置调整逻辑 2024-03-13 12:07:32 +08:00
赵小蒙
86dc22caf3 remove_overlaps_min_spans阈值调整 0.8->0.65 2024-03-13 11:09:37 +08:00
赵小蒙
07abba71c1 draw_bbox工具类逻辑更新 2024-03-13 11:04:36 +08:00
赵小蒙
6396910992 移动modify_y_axis在pipeline中的位置 2024-03-13 10:43:11 +08:00
myhloli
61405b8af8 Merge pull request #3 from myhloli/dev-in-line-bbox
in line equation bbox modification
2024-03-12 19:57:35 +08:00
liukaiwen
6d728ddad0 Merge branch 'master' into dev-in-line-bbox
# Conflicts:
#	magic_pdf/pdf_parse_by_ocr.py
2024-03-12 19:41:04 +08:00
liukaiwen
5b86ca82fe Merge branch 'master' of github.com:myhloli/Magic-PDF 2024-03-12 19:39:08 +08:00
赵小蒙
7512baaaa3 重构drow_bbox为工具类 2024-03-12 19:23:55 +08:00
myhloli
59d817b1bf Merge pull request #2 from icecraft/feat/proc_bbox
Feat/proc bbox
2024-03-12 19:17:09 +08:00
赵小蒙
070139a5bc span重叠删除的阈值0.8->0.5 2024-03-12 19:14:15 +08:00
许瑞
2611e85349 feat: complete self check 2024-03-12 19:14:00 +08:00
许瑞
43581df62b feat: add remove bbox overlap 2024-03-12 19:12:44 +08:00
赵小蒙
61a0c62c46 pdf_info_dict中间态结构调整 2024-03-12 19:06:58 +08:00
liukaiwen
1b4af400cc Merge branch 'master' into dev-in-line-bbox
# Conflicts:
#	demo/draw_bbox.py
2024-03-12 18:44:53 +08:00
liukaiwen
2e32ecfeab Merge branch 'master' of github.com:myhloli/Magic-PDF
# Conflicts:
#	demo/draw_bbox.py
2024-03-12 18:27:42 +08:00
liukaiwen
6b401112e3 add modify inline equation y axis
add false displayed equation to inline equation
2024-03-12 18:27:23 +08:00
赵小蒙
f31117deba debug时自动绘制layout区域和text区域 2024-03-12 18:21:11 +08:00
赵小蒙
ec1a6ef716 增加生成多模态markdown逻辑 2024-03-12 18:20:10 +08:00
赵小蒙
3c8b254574 修复了layout相交的分离算法,并修复layout排序有误的问题 2024-03-12 14:42:13 +08:00
liukaiwen
5513e48ab8 add modify inline equation y axis
add false displayed equation to inline equation
2024-03-12 14:34:15 +08:00
liukaiwen
f9f36c10cf lkw 2024-03-12 14:31:15 +08:00
liukaiwen
514a571836 Merge branch 'master' into dev-in-line-bbox 2024-03-12 11:23:02 +08:00
liukaiwen
1d59509d9d Merge branch 'master' of github.com:myhloli/Magic-PDF 2024-03-12 11:22:39 +08:00
liukaiwen
94a7ba3d23 lkw 2024-03-12 11:22:25 +08:00
赵小蒙
9cc53a5e57 更新清除重叠span的逻辑 2024-03-12 11:21:57 +08:00
liukaiwen
da5091430b Merge branch 'master' into dev-in-line-bbox
# Conflicts:
#	demo/draw_bbox.py
#	demo/ocr_demo.py
#	magic_pdf/pdf_parse_by_ocr.py
#	magic_pdf/pre_proc/ocr_dict_merge.py
2024-03-11 16:52:36 +08:00
liukaiwen
8a860e32e1 lkw 2024-03-11 16:48:36 +08:00
liukaiwen
047a8df2ac lkw 2024-03-11 16:46:10 +08:00
liukaiwen
83deab214e lkw 2024-03-11 16:41:57 +08:00
赵小蒙
72b7894c70 模型输出的layout顺序没有排序,需要后期排序 2024-03-11 14:09:46 +08:00
赵小蒙
012a8f1567 修复多栏模式下lines可能为空导致的异常 2024-03-11 14:08:05 +08:00
赵小蒙
a5f8de9882 ocr模式增加截图功能 2024-03-08 19:32:38 +08:00
liukaiwen
c38c784eaa lkw 2024-03-08 19:08:23 +08:00
赵小蒙
68e83c124f span进入具体的layout后,需要在下次循环前将该span移除 2024-03-08 18:46:31 +08:00
赵小蒙
17b09f7177 ocr pipeline更新 2024-03-08 18:34:19 +08:00
赵小蒙
864e95355f span->line现基于模型的layout进行拼接 2024-03-08 18:24:16 +08:00
赵小蒙
0c279ffccd 更新remove_spans_by_bboxes中选择被删除的span的逻辑 2024-03-08 17:27:46 +08:00
赵小蒙
f9bd004048 对模型的layout坐标转换 2024-03-08 17:13:25 +08:00
赵小蒙
f62d1aa781 将模型和pymu坐标的转换逻辑抽象成方法 2024-03-08 16:50:43 +08:00
赵小蒙
388223f2e0 ocr模式下删除header/page number/footnote/footer 2024-03-08 16:32:37 +08:00
赵小蒙
fcea39d36b 增加ocr模式的layout解析功能 2024-03-07 20:41:41 +08:00
赵小蒙
00f3e329d9 修复一个span可能没有content导致的问题 2024-03-07 16:15:14 +08:00
赵小蒙
caa1588a92 ocr拼接逻辑更新 2024-03-07 16:04:16 +08:00
赵小蒙
a0be4652e6 parse_pdf_by_ocr 逻辑更新 2024-03-06 18:03:55 +08:00
22 changed files with 1116 additions and 229 deletions

View File

@@ -1,20 +0,0 @@
from magic_pdf.libs.commons import fitz # PyMuPDF
# PDF文件路径
pdf_path = "D:\\project\\20231108code-clean\\code-clean\\tmp\\unittest\\download-pdfs\\scihub\\scihub_53700000\\libgen.scimag53724000-53724999.zip_10.1097\\00129191-200509000-00018.pdf"
doc = fitz.open(pdf_path) # Open the PDF
# 你的数据
data = [[[-2, 0, 603, 80, 24]], [[-3, 0, 602, 80, 24]]]
# 对每个页面进行处理
for i, page in enumerate(doc):
# 获取当前页面的数据
page_data = data[i]
for img in page_data:
x0, y0, x1, y1, _ = img
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
doc.save("D:\\project\\20231108code-clean\\code-clean\\tmp\\unittest\\download-pdfs\\scihub\\scihub_53700000\\libgen.scimag53724000-53724999.zip_10.1097\\00129191-200509000-00018_new.pdf")

View File

@@ -1,8 +1,11 @@
import json
import os
from loguru import logger
from pathlib import Path
from magic_pdf.dict2md.ocr_mkcontent import mk_nlp_markdown
from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_nlp_markdown, ocr_mk_mm_markdown
from magic_pdf.libs.commons import join_path
from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
@@ -20,10 +23,46 @@ def save_markdown(markdown_text, input_filepath):
file.write(markdown_text)
if __name__ == '__main__':
ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_0.json"
pdf_info_dict = parse_pdf_by_ocr(ocr_json_file_path)
markdown_text = mk_nlp_markdown(pdf_info_dict)
logger.info(markdown_text)
save_markdown(markdown_text, ocr_json_file_path)
def read_json_file(file_path):
with open(file_path, 'r') as f:
data = json.load(f)
return data
if __name__ == '__main__':
ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.pdf"
ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.json"
# ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.pdf"
# ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.json"
# ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_demo\ocr_1_org.pdf"
# ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_demo\ocr_1_fix.json"
try:
ocr_pdf_model_info = read_json_file(ocr_json_file_path)
pth = Path(ocr_json_file_path)
book_name = pth.name
save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
save_path = join_path(save_tmp_path, "md")
save_path_with_bookname = os.path.join(save_path, book_name)
text_content_save_path = f"{save_path_with_bookname}/book.md"
pdf_info_dict = parse_pdf_by_ocr(
ocr_pdf_path,
None,
ocr_pdf_model_info,
save_path,
book_name,
debug_mode=True)
parent_dir = os.path.dirname(text_content_save_path)
if not os.path.exists(parent_dir):
os.makedirs(parent_dir)
# markdown_content = ocr_mk_nlp_markdown(pdf_info_dict)
markdown_content = ocr_mk_mm_markdown(pdf_info_dict)
with open(text_content_save_path, "w", encoding="utf-8") as f:
f.write(markdown_content)
# logger.info(markdown_content)
# save_markdown(markdown_text, ocr_json_file_path)
except Exception as e:
logger.exception(e)

View File

@@ -1,4 +1,32 @@
def mk_nlp_markdown(pdf_info_dict: dict):
from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
from magic_pdf.libs.ocr_content_type import ContentType
def ocr_mk_nlp_markdown(pdf_info_dict: dict):
markdown = []
for _, page_info in pdf_info_dict.items():
blocks = page_info.get("preproc_blocks")
if not blocks:
continue
for block in blocks:
for line in block['lines']:
line_text = ''
for span in line['spans']:
if not span.get('content'):
continue
content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
if span['type'] == ContentType.InlineEquation:
content = f"${content}$"
elif span['type'] == ContentType.InterlineEquation:
content = f"$$\n{content}\n$$"
line_text += content + ' '
# 在行末添加两个空格以强制换行
markdown.append(line_text.strip() + ' ')
return '\n'.join(markdown)
def ocr_mk_mm_markdown(pdf_info_dict: dict):
markdown = []
@@ -10,12 +38,26 @@ def mk_nlp_markdown(pdf_info_dict: dict):
for line in block['lines']:
line_text = ''
for span in line['spans']:
content = span['content'].replace('$', '\$') # 转义$
if span['type'] == 'inline_equation':
content = f"${content}$"
elif span['type'] == 'displayed_equation':
content = f"$$\n{content}\n$$"
if not span.get('content'):
if not span.get('image_path'):
continue
else:
content = f"![](s3://mllm-raw-media/pdf2md_img/{span['image_path']})"
else:
content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
if span['type'] == ContentType.InlineEquation:
content = f"${content}$"
elif span['type'] == ContentType.InterlineEquation:
content = f"$$\n{content}\n$$"
line_text += content + ' '
# 在行末添加两个空格以强制换行
markdown.append(line_text.strip() + ' ')
return '\n'.join(markdown)
def ocr_mk_mm_standard_format():
'''
content_list
type string image/text/table/equation(行间的单独拿出来行内的和text合并)
'''
pass

View File

@@ -177,6 +177,27 @@ def calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2):
else:
return intersection_area / min_box_area
def calculate_overlap_area_in_bbox1_area_ratio(bbox1, bbox2):
"""
计算box1和box2的重叠面积占bbox1的比例
"""
# Determine the coordinates of the intersection rectangle
x_left = max(bbox1[0], bbox2[0])
y_top = max(bbox1[1], bbox2[1])
x_right = min(bbox1[2], bbox2[2])
y_bottom = min(bbox1[3], bbox2[3])
if x_right < x_left or y_bottom < y_top:
return 0.0
# The area of overlap area
intersection_area = (x_right - x_left) * (y_bottom - y_top)
bbox1_area = (bbox1[2]-bbox1[0])*(bbox1[3]-bbox1[1])
if bbox1_area == 0:
return 0
else:
return intersection_area / bbox1_area
def get_minbox_if_overlap_by_ratio(bbox1, bbox2, ratio):
"""
@@ -189,10 +210,11 @@ def get_minbox_if_overlap_by_ratio(bbox1, bbox2, ratio):
area1 = (x1_max - x1_min) * (y1_max - y1_min)
area2 = (x2_max - x2_min) * (y2_max - y2_min)
overlap_ratio = calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
if overlap_ratio > ratio and area1 < area2:
return bbox1
elif overlap_ratio > ratio and area2 < area1:
return bbox2
if overlap_ratio > ratio:
if area1 <= area2:
return bbox1
else:
return bbox2
else:
return None

View File

@@ -1,4 +1,5 @@
import datetime
import json
import os, re, configparser
import time
@@ -115,6 +116,34 @@ def read_file(pdf_path: str, s3_profile):
with open(pdf_path, "rb") as f:
return f.read()
def get_docx_model_output(pdf_model_output, pdf_model_s3_profile, page_id):
if isinstance(pdf_model_output, str):
model_output_json_path = join_path(pdf_model_output, f"page_{page_id + 1}.json") # 模型输出的页面编号从1开始的
if os.path.exists(model_output_json_path):
json_from_docx = read_file(model_output_json_path, pdf_model_s3_profile)
model_output_json = json.loads(json_from_docx)
else:
try:
model_output_json_path = join_path(pdf_model_output, "model.json")
with open(model_output_json_path, "r", encoding="utf-8") as f:
model_output_json = json.load(f)
model_output_json = model_output_json["doc_layout_result"][page_id]
except:
s3_model_output_json_path = join_path(pdf_model_output, f"page_{page_id + 1}.json")
s3_model_output_json_path = join_path(pdf_model_output, f"{page_id}.json")
#s3_model_output_json_path = join_path(pdf_model_output, f"page_{page_id }.json")
# logger.warning(f"model_output_json_path: {model_output_json_path} not found. try to load from s3: {s3_model_output_json_path}")
s = read_file(s3_model_output_json_path, pdf_model_s3_profile)
return json.loads(s)
elif isinstance(pdf_model_output, list):
model_output_json = pdf_model_output[page_id]
return model_output_json
def list_dir(dir_path:str, s3_profile:str):
"""
列出dir_path下的所有文件

View File

@@ -0,0 +1,9 @@
def get_scale_ratio(ocr_page_info, page):
pix = page.get_pixmap(dpi=72)
pymu_width = int(pix.w)
pymu_height = int(pix.h)
width_from_json = ocr_page_info['page_info']['width']
height_from_json = ocr_page_info['page_info']['height']
horizontal_scale_ratio = width_from_json / pymu_width
vertical_scale_ratio = height_from_json / pymu_height
return horizontal_scale_ratio, vertical_scale_ratio

View File

@@ -0,0 +1,80 @@
from magic_pdf.libs.commons import fitz # PyMuPDF
from magic_pdf.libs.ocr_content_type import ContentType
def draw_bbox_without_number(i, bbox_list, page, rgb_config):
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
rect_coords = fitz.Rect(x0, y0, x1, y1) # Define the rectangle
page.draw_rect(rect_coords, color=new_rgb, fill=None, width=0.5, overlay=True) # Draw the rectangle
def draw_bbox_with_number(i, bbox_list, page, rgb_config):
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
page.draw_rect(rect_coords, color=new_rgb, fill=None, width=0.5, overlay=True) # Draw the rectangle
page.insert_text((x0, y0), str(j + 1), fontsize=10, color=new_rgb) # Insert the index at the top left corner of the rectangle
def draw_layout_bbox(pdf_info_dict, input_path, out_path):
layout_bbox_list = []
dropped_bbox_list = []
for page in pdf_info_dict.values():
page_layout_list = []
page_dropped_list = []
for layout in page['layout_bboxes']:
page_layout_list.append(layout['layout_bbox'])
layout_bbox_list.append(page_layout_list)
for drop_tag, dropped_bboxes in page['droped_bboxes'].items():
for dropped_bbox in dropped_bboxes:
page_dropped_list.append(dropped_bbox)
dropped_bbox_list.append(page_dropped_list)
doc = fitz.open(input_path)
for i, page in enumerate(doc):
draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0])
draw_bbox_without_number(i, dropped_bbox_list, page, [0, 255, 0])
# Save the PDF
doc.save(f"{out_path}/layout.pdf")
def draw_text_bbox(pdf_info_dict, input_path, out_path):
text_list = []
inline_equation_list = []
interline_equation_list = []
for page in pdf_info_dict.values():
page_text_list = []
page_inline_equation_list = []
page_interline_equation_list = []
for block in page['preproc_blocks']:
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'])
text_list.append(page_text_list)
inline_equation_list.append(page_inline_equation_list)
interline_equation_list.append(page_interline_equation_list)
doc = fitz.open(input_path)
for i, page in enumerate(doc):
# 获取当前页面的数据
draw_bbox_without_number(i, text_list, page, [255, 0, 0])
draw_bbox_without_number(i, inline_equation_list, page, [0, 255, 0])
draw_bbox_without_number(i, interline_equation_list, page, [0, 0, 255])
# Save the PDF
doc.save(f"{out_path}/text.pdf")

View File

@@ -18,3 +18,14 @@ def escape_special_markdown_char(pymu_blocks):
span['text'] = span['text'].replace(char, "\\" + char)
return pymu_blocks
def ocr_escape_special_markdown_char(content):
"""
转义正文里对markdown语法有特殊意义的字符
"""
special_chars = ["*", "`", "~", "$"]
for char in special_chars:
content = content.replace(char, "\\" + char)
return content

View File

@@ -0,0 +1,7 @@
class ContentType:
Image = "image"
Table = "table"
Text = "text"
InlineEquation = "inline_equation"
InterlineEquation = "interline_equation"

View File

@@ -1,46 +0,0 @@
from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold
def merge_spans(spans):
# 按照y0坐标排序
spans.sort(key=lambda span: span['bbox'][1])
lines = []
current_line = [spans[0]]
for span in spans[1:]:
# 如果当前的span类型为"displayed_equation" 或者 当前行中已经有"displayed_equation"
if span['type'] == "displayed_equation" or any(s['type'] == "displayed_equation" for s in current_line):
# 则开始新行
lines.append(current_line)
current_line = [span]
continue
# 如果当前的span与当前行的最后一个span在y轴上重叠则添加到当前行
if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
current_line.append(span)
else:
# 否则,开始新行
lines.append(current_line)
current_line = [span]
# 添加最后一行
if current_line:
lines.append(current_line)
# 计算每行的边界框并对每行中的span按照x0进行排序
line_objects = []
for line in lines:
# 按照x0坐标排序
line.sort(key=lambda span: span['bbox'][0])
line_bbox = [
min(span['bbox'][0] for span in line), # x0
min(span['bbox'][1] for span in line), # y0
max(span['bbox'][2] for span in line), # x1
max(span['bbox'][3] for span in line), # y1
]
line_objects.append({
"bbox": line_bbox,
"spans": line,
})
return line_objects

View File

@@ -2,7 +2,7 @@ import time
# from anyio import Path
from magic_pdf.libs.commons import fitz, get_delta_time, get_img_s3_client
from magic_pdf.libs.commons import fitz, get_delta_time, get_img_s3_client, get_docx_model_output
import json
import os
import math
@@ -68,31 +68,6 @@ paraSplitException_msg = ParaSplitException().message
paraMergeException_msg = ParaMergeException().message
def get_docx_model_output(pdf_model_output, pdf_model_s3_profile, page_id):
if isinstance(pdf_model_output, str):
model_output_json_path = join_path(pdf_model_output, f"page_{page_id + 1}.json") # 模型输出的页面编号从1开始的
if os.path.exists(model_output_json_path):
json_from_docx = read_file(model_output_json_path, pdf_model_s3_profile)
model_output_json = json.loads(json_from_docx)
else:
try:
model_output_json_path = join_path(pdf_model_output, "model.json")
with open(model_output_json_path, "r", encoding="utf-8") as f:
model_output_json = json.load(f)
model_output_json = model_output_json["doc_layout_result"][page_id]
except:
s3_model_output_json_path = join_path(pdf_model_output, f"page_{page_id + 1}.json")
s3_model_output_json_path = join_path(pdf_model_output, f"{page_id}.json")
#s3_model_output_json_path = join_path(pdf_model_output, f"page_{page_id }.json")
# logger.warning(f"model_output_json_path: {model_output_json_path} not found. try to load from s3: {s3_model_output_json_path}")
s = read_file(s3_model_output_json_path, pdf_model_s3_profile)
return json.loads(s)
elif isinstance(pdf_model_output, list):
model_output_json = pdf_model_output[page_id]
return model_output_json
def parse_pdf_by_model(

View File

@@ -1,85 +1,261 @@
import json
import os
import time
from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio
from magic_pdf.libs.ocr_dict_merge import merge_spans
from loguru import logger
from magic_pdf.libs.draw_bbox import draw_layout_bbox, draw_text_bbox
from magic_pdf.libs.commons import (
read_file,
join_path,
fitz,
get_img_s3_client,
get_delta_time,
get_docx_model_output,
)
from magic_pdf.libs.coordinate_transform import get_scale_ratio
from magic_pdf.libs.ocr_content_type import ContentType
from magic_pdf.libs.safe_filename import sanitize_filename
from magic_pdf.pre_proc.detect_footer_by_model import parse_footers
from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model
from magic_pdf.pre_proc.detect_header import parse_headers
from magic_pdf.pre_proc.detect_page_number import parse_pageNos
from magic_pdf.pre_proc.ocr_cut_image import cut_image_and_table
from magic_pdf.pre_proc.ocr_detect_layout import layout_detect
from magic_pdf.pre_proc.ocr_dict_merge import (
merge_spans_to_line_by_layout, merge_lines_to_block,
)
from magic_pdf.pre_proc.ocr_span_list_modify import remove_spans_by_bboxes, remove_overlaps_min_spans, \
adjust_bbox_for_standalone_block, modify_y_axis, modify_inline_equation, get_qa_need_list, \
remove_spans_by_bboxes_dict
from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
def read_json_file(file_path):
with open(file_path, 'r') as f:
data = json.load(f)
return data
def construct_page_component(page_id, text_blocks_preproc):
def construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
images, tables, interline_equations, inline_equations,
dropped_text_block, dropped_image_block, dropped_table_block,
need_remove_spans_bboxes_dict):
return_dict = {
'preproc_blocks': text_blocks_preproc,
'page_idx': page_id
'preproc_blocks': blocks,
'layout_bboxes': layout_bboxes,
'page_idx': page_id,
'page_size': [page_w, page_h],
'_layout_tree': layout_tree,
'images': images,
'tables': tables,
'interline_equations': interline_equations,
'inline_equations': inline_equations,
'droped_text_block': dropped_text_block,
'droped_image_block': dropped_image_block,
'droped_table_block': dropped_table_block,
'droped_bboxes': need_remove_spans_bboxes_dict,
}
return return_dict
def parse_pdf_by_ocr(
ocr_json_file_path,
pdf_path,
s3_pdf_profile,
pdf_model_output,
save_path,
book_name,
pdf_model_profile=None,
image_s3_config=None,
start_page_id=0,
end_page_id=None,
debug_mode=False,
):
ocr_pdf_info = read_json_file(ocr_json_file_path)
pdf_bytes = read_file(pdf_path, s3_pdf_profile)
save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
book_name = sanitize_filename(book_name)
md_bookname_save_path = ""
if debug_mode:
save_path = join_path(save_tmp_path, "md")
pdf_local_path = join_path(save_tmp_path, "download-pdfs", book_name)
if not os.path.exists(os.path.dirname(pdf_local_path)):
# 如果目录不存在,创建它
os.makedirs(os.path.dirname(pdf_local_path))
md_bookname_save_path = join_path(save_tmp_path, "md", book_name)
if not os.path.exists(md_bookname_save_path):
# 如果目录不存在,创建它
os.makedirs(md_bookname_save_path)
with open(pdf_local_path + ".pdf", "wb") as pdf_file:
pdf_file.write(pdf_bytes)
pdf_docs = fitz.open("pdf", pdf_bytes)
# 初始化空的pdf_info_dict
pdf_info_dict = {}
end_page_id = end_page_id if end_page_id else len(ocr_pdf_info) - 1
img_s3_client = get_img_s3_client(save_path, image_s3_config)
start_time = time.time()
end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
for page_id in range(start_page_id, end_page_id + 1):
ocr_page_info = ocr_pdf_info[page_id]
layout_dets = ocr_page_info['layout_dets']
# 获取当前页的page对象
page = pdf_docs[page_id]
# 获取当前页的宽高
page_w = page.rect.width
page_h = page.rect.height
if debug_mode:
time_now = time.time()
logger.info(
f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
)
start_time = time_now
# 获取当前页的模型数据
ocr_page_info = get_docx_model_output(
pdf_model_output, pdf_model_profile, page_id
)
"""从json中获取每页的页码、页眉、页脚的bbox"""
page_no_bboxes = parse_pageNos(page_id, page, ocr_page_info)
header_bboxes = parse_headers(page_id, page, ocr_page_info)
footer_bboxes = parse_footers(page_id, page, ocr_page_info)
footnote_bboxes = parse_footnotes_by_model(
page_id, page, ocr_page_info, md_bookname_save_path, debug_mode=debug_mode
)
# 构建需要remove的bbox列表
# need_remove_spans_bboxes = []
# need_remove_spans_bboxes.extend(page_no_bboxes)
# need_remove_spans_bboxes.extend(header_bboxes)
# need_remove_spans_bboxes.extend(footer_bboxes)
# need_remove_spans_bboxes.extend(footnote_bboxes)
# 构建需要remove的bbox字典
need_remove_spans_bboxes_dict = {
"page_no": page_no_bboxes,
"header": header_bboxes,
"footer": footer_bboxes,
"footnote": footnote_bboxes,
}
layout_dets = ocr_page_info["layout_dets"]
spans = []
# 计算模型坐标和pymu坐标的缩放比例
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
ocr_page_info, page
)
for layout_det in layout_dets:
category_id = layout_det['category_id']
allow_category_id_list = [13, 14, 15]
category_id = layout_det["category_id"]
allow_category_id_list = [1, 7, 13, 14, 15]
if category_id in allow_category_id_list:
x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
bbox = [int(x0), int(y0), int(x1), int(y1)]
# 13: 'embedding', # 嵌入公式
# 14: 'isolated', # 单行公式
# 15: 'ocr_text', # ocr识别文本
x0, y0, _, _, x1, y1, _, _ = layout_det["poly"]
bbox = [
int(x0 / horizontal_scale_ratio),
int(y0 / vertical_scale_ratio),
int(x1 / horizontal_scale_ratio),
int(y1 / vertical_scale_ratio),
]
# 删除高度或者宽度为0的spans
if bbox[2] - bbox[0] == 0 or bbox[3] - bbox[1] == 0:
continue
"""要删除的"""
# 3: 'header', # 页眉
# 4: 'page number', # 页码
# 5: 'footnote', # 脚注
# 6: 'footer', # 页脚
"""当成span拼接的"""
# 1: 'image', # 图片
# 7: 'table', # 表格
# 13: 'inline_equation', # 行内公式
# 14: 'interline_equation', # 行间公式
# 15: 'text', # ocr识别文本
"""layout信息"""
# 11: 'full column', # 单栏
# 12: 'sub column', # 多栏
span = {
'bbox': bbox,
"bbox": bbox,
}
if category_id == 13:
span['content'] = layout_det['latex']
span['type'] = 'inline_equation'
if category_id == 1:
span["type"] = ContentType.Image
elif category_id == 7:
span["type"] = ContentType.Table
elif category_id == 13:
span["content"] = layout_det["latex"]
span["type"] = ContentType.InlineEquation
elif category_id == 14:
span['content'] = layout_det['latex']
span['type'] = 'displayed_equation'
span["content"] = layout_det["latex"]
span["type"] = ContentType.InterlineEquation
elif category_id == 15:
span['content'] = layout_det['text']
span['type'] = 'text'
span["content"] = layout_det["text"]
span["type"] = ContentType.Text
# print(span)
spans.append(span)
else:
continue
# 合并重叠的spans
for span1 in spans.copy():
for span2 in spans.copy():
if span1 != span2:
overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.8)
if overlap_box is not None:
bbox_to_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
if bbox_to_remove is not None:
spans.remove(bbox_to_remove)
# 将spans合并成line
lines = merge_spans(spans)
# 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
blocks = []
for line in lines:
blocks.append({
"bbox": line['bbox'],
"lines": [line],
})
# 删除重叠spans中较小的那些
spans = remove_overlaps_min_spans(spans)
# 删除remove_span_block_bboxes中的bbox
# spans = remove_spans_by_bboxes(spans, need_remove_spans_bboxes)
# 按qa要求增加drop相关数据
spans, dropped_text_block, dropped_image_block, dropped_table_block = remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict)
# 对image和table截图
spans = cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client)
# 行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)
displayed_list = []
text_inline_lines = []
modify_y_axis(spans, displayed_list, text_inline_lines)
# 模型识别错误的行间公式, type类型转换成行内公式
spans = modify_inline_equation(spans, displayed_list, text_inline_lines)
# bbox去除粘连
spans = remove_overlap_between_bbox(spans)
# 对tpye=["interline_equation", "image", "table"]进行额外处理,如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
spans = adjust_bbox_for_standalone_block(spans)
# 从ocr_page_info中解析layout信息(按自然阅读方向排序,并修复重叠和交错的bad case)
layout_bboxes, layout_tree = layout_detect(ocr_page_info['subfield_dets'], page, ocr_page_info)
# 将spans合并成line(在layout内,从上到下,从左到右)
lines = merge_spans_to_line_by_layout(spans, layout_bboxes)
# 将lines合并成block
blocks = merge_lines_to_block(lines)
# 根据block合并段落
# 获取QA需要外置的list
images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
# 构造pdf_info_dict
page_info = construct_page_component(page_id, blocks)
page_info = construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
images, tables, interline_equations, inline_equations,
dropped_text_block, dropped_image_block, dropped_table_block,
need_remove_spans_bboxes_dict)
pdf_info_dict[f"page_{page_id}"] = page_info
return pdf_info_dict
# 在测试时,保存调试信息
if debug_mode:
params_file_save_path = join_path(
save_tmp_path, "md", book_name, "preproc_out.json"
)
with open(params_file_save_path, "w", encoding="utf-8") as f:
json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
# drow_bbox
draw_layout_bbox(pdf_info_dict, pdf_path, md_bookname_save_path)
draw_text_bbox(pdf_info_dict, pdf_path, md_bookname_save_path)
return pdf_info_dict

View File

@@ -3,6 +3,7 @@ import sys
import time
from urllib.parse import quote
from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_nlp_markdown, ocr_mk_mm_markdown
from magic_pdf.libs.commons import read_file, join_path, parse_bucket_key, formatted_time
from magic_pdf.libs.drop_reason import DropReason
from magic_pdf.libs.json_compressor import JsonCompressor
@@ -13,6 +14,7 @@ from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
from loguru import logger
from app.common.s3 import get_s3_config, get_s3_client
from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
def exception_handler(jso: dict, e):
@@ -23,6 +25,27 @@ def exception_handler(jso: dict, e):
return jso
def get_data_type(jso: dict):
data_type = jso.get('data_type')
if data_type is None:
data_type = jso.get('file_type')
return data_type
def get_bookid(jso: dict):
book_id = jso.get('bookid')
if book_id is None:
book_id = jso.get('original_file_id')
return book_id
def get_data_source(jso: dict):
data_source = jso.get('data_source')
if data_source is None:
data_source = jso.get('file_source')
return data_source
def meta_scan(jso: dict, doc_layout_check=True) -> dict:
s3_pdf_path = jso.get('file_location')
s3_config = get_s3_config(s3_pdf_path)
@@ -32,7 +55,7 @@ def meta_scan(jso: dict, doc_layout_check=True) -> dict:
jso['drop_reason'] = DropReason.MISS_DOC_LAYOUT_RESULT
return jso
try:
data_source = jso.get('data_source')
data_source = get_data_source(jso)
file_id = jso.get('file_id')
book_name = data_source + "/" + file_id
@@ -78,7 +101,7 @@ def classify_by_type(jso: dict, debug_mode=False) -> dict:
# 开始正式逻辑
try:
pdf_meta = jso.get('pdf_meta')
data_source = jso.get('data_source')
data_source = get_data_source(jso)
file_id = jso.get('file_id')
book_name = data_source + "/" + file_id
total_page = pdf_meta["total_page"]
@@ -140,11 +163,11 @@ def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:
pass
else:# 如果debug没开则检测是否有needdrop字段
if jso.get('need_drop', False):
logger.info(f"book_name is:{jso['data_source']}/{jso['file_id']} need drop", file=sys.stderr)
logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop", file=sys.stderr)
jso["dropped"] = True
return jso
try:
data_source = jso.get('data_source')
data_source = get_data_source(jso)
file_id = jso.get('file_id')
book_name = data_source + "/" + file_id
title = jso.get('title')
@@ -195,7 +218,7 @@ def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:
def drop_needdrop_pdf(jso: dict) -> dict:
if jso.get('need_drop', False):
logger.info(f"book_name is:{jso['data_source']}/{jso['file_id']} need drop", file=sys.stderr)
logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop", file=sys.stderr)
jso["dropped"] = True
return jso
@@ -206,7 +229,7 @@ def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
pass
else:# 如果debug没开则检测是否有needdrop字段
if jso.get('need_drop', False):
book_name = join_path(jso['data_source'], jso['file_id'])
book_name = join_path(get_data_source(jso), jso['file_id'])
logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
jso["dropped"] = True
return jso
@@ -216,7 +239,7 @@ def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
markdown_content = mk_nlp_markdown(pdf_intermediate_dict)
jso["content"] = markdown_content
logger.info(f"book_name is:{jso['data_source']}/{jso['file_id']},markdown content length is {len(markdown_content)}", file=sys.stderr)
logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}", file=sys.stderr)
# 把无用的信息清空
jso["doc_layout_result"] = ""
jso["pdf_intermediate_dict"] = ""
@@ -237,7 +260,7 @@ def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
s3_pdf_path = jso.get('file_location')
s3_config = get_s3_config(s3_pdf_path)
model_output_json_list = jso.get('doc_layout_result')
data_source = jso.get('data_source')
data_source = get_data_source(jso)
file_id = jso.get('file_id')
book_name = data_source + "/" + file_id
@@ -290,5 +313,79 @@ def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
return jso
def ocr_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
# 检测debug开关
if debug_mode:
pass
else: # 如果debug没开则检测是否有needdrop字段
if jso.get('need_drop', False):
return jso
s3_pdf_path = jso.get('file_location')
s3_config = get_s3_config(s3_pdf_path)
model_output_json_list = jso.get('doc_layout_result')
data_source = get_data_source(jso)
file_id = jso.get('file_id')
book_name = data_source + "/" + file_id
try:
save_path = "s3://mllm-raw-media/pdf2md_img/"
image_s3_config = get_s3_config(save_path)
start_time = time.time() # 记录开始时间
# 先打印一下book_name和解析开始的时间
logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
pdf_info_dict = parse_pdf_by_ocr(
s3_pdf_path,
s3_config,
model_output_json_list,
save_path,
book_name,
pdf_model_profile=None,
image_s3_config=image_s3_config,
start_page_id=start_page_id,
debug_mode=debug_mode
)
if pdf_info_dict.get('need_drop', False): # 如果返回的字典里有need_drop则提取drop_reason并跳过本次解析
jso['need_drop'] = True
jso['drop_reason'] = pdf_info_dict["drop_reason"]
else: # 正常返回,将 pdf_info_dict 压缩并存储
pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
jso['pdf_intermediate_dict'] = pdf_info_dict
end_time = time.time() # 记录完成时间
parse_time = int(end_time - start_time) # 计算执行时间
# 解析完成后打印一下book_name和耗时
logger.info(f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
file=sys.stderr)
jso['parse_time'] = parse_time
except Exception as e:
jso = exception_handler(jso, e)
return jso
def ocr_pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
if debug_mode:
pass
else: # 如果debug没开则检测是否有needdrop字段
if jso.get('need_drop', False):
book_name = join_path(get_data_source(jso), jso['file_id'])
logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
jso["dropped"] = True
return jso
try:
pdf_intermediate_dict = jso['pdf_intermediate_dict']
# 将 pdf_intermediate_dict 解压
pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
markdown_content = ocr_mk_mm_markdown(pdf_intermediate_dict)
jso["content"] = markdown_content
logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}", file=sys.stderr)
# 把无用的信息清空
jso["doc_layout_result"] = ""
jso["pdf_intermediate_dict"] = ""
jso["pdf_meta"] = ""
except Exception as e:
jso = exception_handler(jso, e)
return jso
if __name__ == "__main__":
pass

View File

@@ -1,4 +1,5 @@
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_footers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
@@ -8,23 +9,12 @@ def parse_footers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
:param res_dir_path: str类型是每一个pdf文档在当前.py文件的目录下生成一个与pdf文档同名的文件夹res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型把pdf文档送入DocXChain模型中后提取bbox结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#--------- 通过json_from_DocXchain来获取 footer ---------#
footer_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
@@ -42,10 +32,10 @@ def parse_footers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面artBox偏移了。不在0,0
# R += pageL
# U += pageU

View File

@@ -1,5 +1,6 @@
from collections import Counter
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_footnotes_by_model(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict, md_bookname_save_path, debug_mode=False):
@@ -9,22 +10,12 @@ def parse_footnotes_by_model(page_ID: int, page: fitz.Page, json_from_DocXchain_
:param res_dir_path: str类型是每一个pdf文档在当前.py文件的目录下生成一个与pdf文档同名的文件夹res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型把pdf文档送入DocXChain模型中后提取bbox结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#--------- 通过json_from_DocXchain来获取 footnote ---------#
footnote_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
@@ -42,10 +33,10 @@ def parse_footnotes_by_model(page_ID: int, page: fitz.Page, json_from_DocXchain_
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面artBox偏移了。不在0,0
# R += pageL
# U += pageU

View File

@@ -1,4 +1,5 @@
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_headers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
@@ -8,22 +9,12 @@ def parse_headers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
:param res_dir_path: str类型是每一个pdf文档在当前.py文件的目录下生成一个与pdf文档同名的文件夹res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型把pdf文档送入DocXChain模型中后提取bbox结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#--------- 通过json_from_DocXchain来获取 header ---------#
header_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
@@ -41,10 +32,10 @@ def parse_headers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面artBox偏移了。不在0,0
# R += pageL
# U += pageU

View File

@@ -1,4 +1,5 @@
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_pageNos(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
@@ -8,22 +9,12 @@ def parse_pageNos(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
:param res_dir_path: str类型是每一个pdf文档在当前.py文件的目录下生成一个与pdf文档同名的文件夹res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型把pdf文档送入DocXChain模型中后提取bbox结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#--------- 通过json_from_DocXchain来获取 pageNo ---------#
pageNo_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
@@ -41,10 +32,10 @@ def parse_pageNos(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面artBox偏移了。不在0,0
# R += pageL
# U += pageU

View File

@@ -0,0 +1,20 @@
from magic_pdf.libs.commons import join_path
from magic_pdf.libs.ocr_content_type import ContentType
from magic_pdf.libs.pdf_image_tools import cut_image
def cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client):
def s3_return_path(type):
return join_path(book_name, type)
def img_save_path(type):
return join_path(save_path, s3_return_path(type))
for span in spans:
span_type = span['type']
if span_type == ContentType.Image:
span['image_path'] = cut_image(span['bbox'], page_id, page, img_save_path('images'), s3_return_path=s3_return_path('images'), img_s3_client=img_s3_client)
elif span_type == ContentType.Table:
span['image_path'] = cut_image(span['bbox'], page_id, page, img_save_path('tables'), s3_return_path=s3_return_path('tables'), img_s3_client=img_s3_client)
return spans

View File

@@ -0,0 +1,133 @@
import fitz
from magic_pdf.layout.layout_sort import get_bboxes_layout
from magic_pdf.libs.boxbase import _is_part_overlap, _is_in
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def get_center_point(bbox):
"""
根据边界框坐标信息,计算出该边界框的中心点坐标。
Args:
bbox (list): 边界框坐标信息包含四个元素分别为左上角x坐标、左上角y坐标、右下角x坐标、右下角y坐标。
Returns:
list: 中心点坐标信息包含两个元素分别为x坐标和y坐标。
"""
return [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
def get_area(bbox):
"""
根据边界框坐标信息,计算出该边界框的面积。
Args:
bbox (list): 边界框坐标信息包含四个元素分别为左上角x坐标、左上角y坐标、右下角x坐标、右下角y坐标。
Returns:
float: 该边界框的面积。
"""
return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
def adjust_layouts(layout_bboxes, page_boundry, page_id):
# 遍历所有布局框
for i in range(len(layout_bboxes)):
# 遍历当前布局框之后的布局框
for j in range(i + 1, len(layout_bboxes)):
# 判断两个布局框是否重叠
if _is_part_overlap(layout_bboxes[i], layout_bboxes[j]):
# 计算每个布局框的中心点坐标和面积
area_i = get_area(layout_bboxes[i])
area_j = get_area(layout_bboxes[j])
# 较大布局框和较小布局框的赋值
if area_i > area_j:
larger_layout, smaller_layout = layout_bboxes[i], layout_bboxes[j]
else:
larger_layout, smaller_layout = layout_bboxes[j], layout_bboxes[i]
center_large = get_center_point(larger_layout)
center_small = get_center_point(smaller_layout)
# 计算横向和纵向的距离差
distance_x = center_large[0] - center_small[0]
distance_y = center_large[1] - center_small[1]
# 根据距离差判断重叠方向并修正边界
if abs(distance_x) > abs(distance_y): # 左右重叠
if distance_x > 0 and larger_layout[0] < smaller_layout[2]:
larger_layout[0] = smaller_layout[2]+1
if distance_x < 0 and larger_layout[2] > smaller_layout[0]:
larger_layout[2] = smaller_layout[0]-1
else: # 上下重叠
if distance_y > 0 and larger_layout[1] < smaller_layout[3]:
larger_layout[1] = smaller_layout[3]+1
if distance_y < 0 and larger_layout[3] > smaller_layout[1]:
larger_layout[3] = smaller_layout[1]-1
# 排序调整布局边界框列表
new_bboxes = []
for layout_bbox in layout_bboxes:
new_bboxes.append([layout_bbox[0], layout_bbox[1], layout_bbox[2], layout_bbox[3], None, None, None, None, None, None, None, None, None])
layout_bboxes, layout_tree = get_bboxes_layout(new_bboxes, page_boundry, page_id)
# 返回排序调整后的布局边界框列表
return layout_bboxes, layout_tree
def layout_detect(layout_info, page: fitz.Page, ocr_page_info):
"""
对输入的布局信息进行解析,提取出每个子布局的边界框,并对所有子布局进行排序调整。
Args:
layout_info (list): 包含子布局信息的列表,每个子布局信息为字典类型,包含'poly'字段,表示子布局的边界框坐标信息。
Returns:
list: 经过排序调整后的所有子布局边界框信息的列表,每个边界框信息为字典类型,包含'layout_bbox'字段,表示边界框的坐标信息。
"""
page_id = ocr_page_info['page_info']['page_no']-1
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(ocr_page_info, page)
# 初始化布局边界框列表
layout_bboxes = []
# 遍历每个子布局
for sub_layout in layout_info:
# 提取子布局的边界框坐标信息
x0, y0, _, _, x1, y1, _, _ = sub_layout['poly']
bbox = [int(x0 / horizontal_scale_ratio), int(y0 / vertical_scale_ratio),
int(x1 / horizontal_scale_ratio), int(y1 / vertical_scale_ratio)]
# 将子布局的边界框添加到列表中
layout_bboxes.append(bbox)
# 初始化新的布局边界框列表
new_layout_bboxes = []
# 遍历每个布局边界框
for i in range(len(layout_bboxes)):
# 初始化标记变量,用于判断当前边界框是否需要保留
keep = True
# 获取当前边界框的坐标信息
box_i = layout_bboxes[i]
# 遍历其他边界框
for j in range(len(layout_bboxes)):
# 排除当前边界框自身
if i != j:
# 获取其他边界框的坐标信息
box_j = layout_bboxes[j]
# 检测box_i是否被box_j包含
if _is_in(box_i, box_j):
# 如果当前边界框被其他边界框包含,则标记为不需要保留
keep = False
# 跳出内层循环
break
# 如果当前边界框需要保留,则添加到新的布局边界框列表中
if keep:
new_layout_bboxes.append(layout_bboxes[i])
# 对新的布局边界框列表进行排序调整
page_width = page.rect.width
page_height = page.rect.height
page_boundry = [0, 0, page_width, page_height]
layout_bboxes, layout_tree = adjust_layouts(new_layout_bboxes, page_boundry, page_id)
# 返回排序调整后的布局边界框列表
return layout_bboxes, layout_tree

View File

@@ -0,0 +1,100 @@
from loguru import logger
from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold, get_minbox_if_overlap_by_ratio, \
calculate_overlap_area_in_bbox1_area_ratio
from magic_pdf.libs.ocr_content_type import ContentType
# 将每一个line中的span从左到右排序
def line_sort_spans_by_left_to_right(lines):
line_objects = []
for line in lines:
# 按照x0坐标排序
line.sort(key=lambda span: span['bbox'][0])
line_bbox = [
min(span['bbox'][0] for span in line), # x0
min(span['bbox'][1] for span in line), # y0
max(span['bbox'][2] for span in line), # x1
max(span['bbox'][3] for span in line), # y1
]
line_objects.append({
"bbox": line_bbox,
"spans": line,
})
return line_objects
def merge_spans_to_line(spans):
# 按照y0坐标排序
spans.sort(key=lambda span: span['bbox'][1])
lines = []
current_line = [spans[0]]
for span in spans[1:]:
# 如果当前的span类型为"interline_equation" 或者 当前行中已经有"interline_equation"
# image和table类型同上
if span['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] or any(
s['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] for s in current_line):
# 则开始新行
lines.append(current_line)
current_line = [span]
continue
# 如果当前的span与当前行的最后一个span在y轴上重叠则添加到当前行
if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
current_line.append(span)
else:
# 否则,开始新行
lines.append(current_line)
current_line = [span]
# 添加最后一行
if current_line:
lines.append(current_line)
return lines
def merge_spans_to_line_by_layout(spans, layout_bboxes):
lines = []
new_spans = []
for item in layout_bboxes:
layout_bbox = item['layout_bbox']
# 遍历spans,将每个span放入对应的layout中
layout_sapns = []
for span in spans:
if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], layout_bbox) > 0.65:
layout_sapns.append(span)
# 如果layout_sapns不为空则放入new_spans中
if len(layout_sapns) > 0:
new_spans.append(layout_sapns)
# 从spans删除已经放入layout_sapns中的span
for layout_sapn in layout_sapns:
spans.remove(layout_sapn)
if len(new_spans) > 0:
for layout_sapns in new_spans:
layout_lines = merge_spans_to_line(layout_sapns)
lines.extend(layout_lines)
#对line中的span进行排序
lines = line_sort_spans_by_left_to_right(lines)
return lines
def merge_lines_to_block(lines):
# 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
blocks = []
for line in lines:
blocks.append(
{
"bbox": line["bbox"],
"lines": [line],
}
)
return blocks

View File

@@ -0,0 +1,207 @@
from loguru import logger
from magic_pdf.libs.boxbase import calculate_overlap_area_in_bbox1_area_ratio, get_minbox_if_overlap_by_ratio, \
__is_overlaps_y_exceeds_threshold
from magic_pdf.libs.ocr_content_type import ContentType
def remove_overlaps_min_spans(spans):
# 删除重叠spans中较小的那些
for span1 in spans.copy():
for span2 in spans.copy():
if span1 != span2:
overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.65)
if overlap_box is not None:
bbox_to_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
if bbox_to_remove is not None:
spans.remove(bbox_to_remove)
return spans
def remove_spans_by_bboxes(spans, need_remove_spans_bboxes):
# 遍历spans, 判断是否在removed_span_block_bboxes中
# 如果是, 则删除该span 否则, 保留该span
need_remove_spans = []
for span in spans:
for removed_bbox in need_remove_spans_bboxes:
if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], removed_bbox) > 0.5:
need_remove_spans.append(span)
break
for span in need_remove_spans:
spans.remove(span)
return spans
def remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict):
dropped_text_block = []
dropped_image_block = []
dropped_table_block = []
for drop_tag, removed_bboxes in need_remove_spans_bboxes_dict.items():
# logger.info(f"remove spans by bbox dict, drop_tag: {drop_tag}, removed_bboxes: {removed_bboxes}")
need_remove_spans = []
for span in spans:
for removed_bbox in removed_bboxes:
if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], removed_bbox) > 0.5:
need_remove_spans.append(span)
break
for span in need_remove_spans:
spans.remove(span)
span['tag'] = drop_tag
if span['type'] in [ContentType.Text, ContentType.InlineEquation, ContentType.InterlineEquation]:
dropped_text_block.append(span)
elif span['type'] == ContentType.Image:
dropped_image_block.append(span)
elif span['type'] == ContentType.Table:
dropped_table_block.append(span)
return spans, dropped_text_block, dropped_image_block, dropped_table_block
def adjust_bbox_for_standalone_block(spans):
# 对tpye=["interline_equation", "image", "table"]进行额外处理,如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
for sb_span in spans:
if sb_span['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table]:
for text_span in spans:
if text_span['type'] in [ContentType.Text, ContentType.InlineEquation]:
# 判断span2的纵向高度是否被span所覆盖
if sb_span['bbox'][1] < text_span['bbox'][1] and sb_span['bbox'][3] > text_span['bbox'][3]:
# 判断span2是否在span左边
if text_span['bbox'][0] < sb_span['bbox'][0]:
# 调整span的y0和span2的y0一致
sb_span['bbox'][1] = text_span['bbox'][1]
return spans
def modify_y_axis(spans: list, displayed_list: list, text_inline_lines: list):
# displayed_list = []
spans.sort(key=lambda span: span['bbox'][1])
lines = []
current_line = [spans[0]]
if spans[0]["type"] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table]:
displayed_list.append(spans[0])
line_first_y0 = spans[0]["bbox"][1]
line_first_y = spans[0]["bbox"][3]
# 用于给行间公式搜索
# text_inline_lines = []
for span in spans[1:]:
# if span.get("content","") == "78.":
# print("debug")
# 如果当前的span类型为"interline_equation" 或者 当前行中已经有"interline_equation"
# image和table类型同上
if span['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] or any(
s['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] for s in current_line):
# 传入
if span["type"] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table]:
displayed_list.append(span)
# 则开始新行
lines.append(current_line)
if len(current_line) > 1 or current_line[0]["type"] in [ContentType.Text, ContentType.InlineEquation]:
text_inline_lines.append((current_line, (line_first_y0, line_first_y)))
current_line = [span]
line_first_y0 = span["bbox"][1]
line_first_y = span["bbox"][3]
continue
# 如果当前的span与当前行的最后一个span在y轴上重叠则添加到当前行
if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
if span["type"] == "text":
line_first_y0 = span["bbox"][1]
line_first_y = span["bbox"][3]
current_line.append(span)
else:
# 否则,开始新行
lines.append(current_line)
text_inline_lines.append((current_line, (line_first_y0, line_first_y)))
current_line = [span]
line_first_y0 = span["bbox"][1]
line_first_y = span["bbox"][3]
# 添加最后一行
if current_line:
lines.append(current_line)
if len(current_line) > 1 or current_line[0]["type"] in [ContentType.Text, ContentType.InlineEquation]:
text_inline_lines.append((current_line, (line_first_y0, line_first_y)))
for line in text_inline_lines:
# 按照x0坐标排序
current_line = line[0]
current_line.sort(key=lambda span: span['bbox'][0])
# 调整每一个文字行内bbox统一
for line in text_inline_lines:
current_line, (line_first_y0, line_first_y) = line
for span in current_line:
span["bbox"][1] = line_first_y0
span["bbox"][3] = line_first_y
# return spans, displayed_list, text_inline_lines
def modify_inline_equation(spans: list, displayed_list: list, text_inline_lines: list):
# 错误行间公式转行内公式
j = 0
for i in range(len(displayed_list)):
# if i == 8:
# print("debug")
span = displayed_list[i]
span_y0, span_y = span["bbox"][1], span["bbox"][3]
while j < len(text_inline_lines):
text_line = text_inline_lines[j]
y0, y1 = text_line[1]
if (
span_y0 < y0 and span_y > y0 or span_y0 < y1 and span_y > y1 or span_y0 < y0 and span_y > y1) and __is_overlaps_y_exceeds_threshold(
span['bbox'], (0, y0, 0, y1)):
# 调整公式类型
if span["type"] == ContentType.InterlineEquation:
# 最后一行是行间公式
if j + 1 >= len(text_inline_lines):
span["type"] = ContentType.InlineEquation
span["bbox"][1] = y0
span["bbox"][3] = y1
else:
# 行间公式旁边有多行文字或者行间公式比文字高3倍则不转换
y0_next, y1_next = text_inline_lines[j + 1][1]
if not __is_overlaps_y_exceeds_threshold(span['bbox'], (0, y0_next, 0, y1_next)) and 3 * (
y1 - y0) > span_y - span_y0:
span["type"] = ContentType.InlineEquation
span["bbox"][1] = y0
span["bbox"][3] = y1
break
elif span_y < y0 or span_y0 < y0 and span_y > y0 and not __is_overlaps_y_exceeds_threshold(span['bbox'],
(0, y0, 0, y1)):
break
else:
j += 1
return spans
def get_qa_need_list(blocks):
# 创建 images, tables, interline_equations, inline_equations 的副本
images = []
tables = []
interline_equations = []
inline_equations = []
for block in blocks:
for line in block["lines"]:
for span in line["spans"]:
if span["type"] == ContentType.Image:
images.append(span)
elif span["type"] == ContentType.Table:
tables.append(span)
elif span["type"] == ContentType.InlineEquation:
inline_equations.append(span)
elif span["type"] == ContentType.InterlineEquation:
interline_equations.append(span)
else:
continue
return images, tables, interline_equations, inline_equations

View File

@@ -0,0 +1,43 @@
from magic_pdf.libs.boxbase import _is_in_or_part_overlap, _is_in
def _remove_overlap_between_bbox(spans):
res = []
for v in spans:
for i in range(len(res)):
if _is_in(res[i]["bbox"], v["bbox"]):
continue
if _is_in_or_part_overlap(res[i]["bbox"], v["bbox"]):
ix0, iy0, ix1, iy1 = res[i]["bbox"]
x0, y0, x1, y1 = v["bbox"]
diff_x = min(x1, ix1) - max(x0, ix0)
diff_y = min(y1, iy1) - max(y0, iy0)
if diff_y > diff_x:
if x1 >= ix1:
mid = (x0 + ix1) // 2
ix1 = min(mid, ix1)
x0 = max(mid + 1, x0)
else:
mid = (ix0 + x1) // 2
ix0 = max(mid + 1, ix0)
x1 = min(mid, x1)
else:
if y1 >= iy1:
mid = (y0 + iy1) // 2
y0 = max(mid + 1, y0)
iy1 = min(iy1, mid)
else:
mid = (iy0 + y1) // 2
y1 = min(y1, mid)
iy0 = max(mid + 1, iy0)
res[i]["bbox"] = [ix0, iy0, ix1, iy1]
v["bbox"] = [x0, y0, x1, y1]
res.append(v)
return res
def remove_overlap_between_bbox(spans):
return _remove_overlap_between_bbox(spans)