修复spans为空list导致的IndexError: list index out of range

# Conflicts:
#	magic_pdf/dict2md/ocr_mkcontent.py

.gitignore

@@ -31,5 +31,6 @@ tmp
 .vscode
 .vscode/
 /tests/
+ocr_demo
 
 /app/common/__init__.py

demo/draw_bbox.py

@@ -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")

demo/ocr_demo.py

@@ -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 mk_mm_markdown2, mk_nlp_markdown, 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,47 @@ 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"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.pdf"
+    ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.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 = mk_nlp_markdown(pdf_info_dict)
+        markdown_content = mk_mm_markdown2(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)

magic_pdf/dict2md/ocr_mkcontent.py

@@ -1,4 +1,33 @@
-def mk_nlp_markdown(pdf_info_dict: dict):
+from magic_pdf.libs.commons import s3_image_save_path, join_path
+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 +39,51 @@ 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"![]({join_path(s3_image_save_path, 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 mk_mm_markdown2(pdf_info_dict:dict):
+    markdown = []
+    for _, page_info in pdf_info_dict.items():
+        paras = page_info.get("para_blocks")
+        if not paras:
+            continue
+        for para in paras:
+            para_text = ''
+            for line in para:
+                for span in line['spans']:
+                    span_type = span.get('type')
+                    if span_type == ContentType.Text:
+                        para_text += span['content']
+                    elif span_type == ContentType.InlineEquation:
+                        para_text += f" ${span['content']}$ "
+                    elif span_type == ContentType.InterlineEquation:
+                        para_text += f"$$\n{span['content']}\n$$ "
+                    elif span_type == ContentType.Image:
+                        para_text += f"![]({join_path(s3_image_save_path, span['image_path'])})"
+            markdown.append(para_text)
+
+    return '\n\n'.join(markdown)
+
+
+def ocr_mk_mm_standard_format():
+    '''
+    content_list
+    type string image/text/table/equation(行间的单独拿出来，行内的和text合并)
+
+    '''
+    pass

magic_pdf/libs/boxbase.py

@@ -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
 

magic_pdf/libs/commons.py

@@ -1,4 +1,5 @@
 import datetime
+import json
 import os, re, configparser
 import time
 
@@ -15,13 +16,16 @@ def get_delta_time(input_time):
 
 
 def join_path(*args):
-    return '/'.join(s.rstrip('/') for s in args)
+    return '/'.join(str(s).rstrip('/') for s in args)
+
 
 #配置全局的errlog_path，方便demo同步引用
 error_log_path = "s3://llm-pdf-text/err_logs/"
 # json_dump_path = "s3://pdf_books_temp/json_dump/" # 这条路径仅用于临时本地测试,不能提交到main
 json_dump_path = "s3://llm-pdf-text/json_dump/"
 
+s3_image_save_path = "s3://mllm-raw-media/pdf2md_img/"
+
 
 def get_top_percent_list(num_list, percent):
     """
@@ -115,6 +119,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下的所有文件

magic_pdf/libs/coordinate_transform.py

@@ -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

magic_pdf/libs/draw_bbox.py

@@ -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")

magic_pdf/libs/markdown_utils.py

@@ -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

magic_pdf/libs/ocr_content_type.py

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

magic_pdf/libs/ocr_dict_merge.py

@@ -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

magic_pdf/para/para_split.py

@@ -0,0 +1,206 @@
+from sklearn.cluster import DBSCAN
+import numpy as np
+from loguru import logger
+
+from magic_pdf.libs.boxbase import _is_in
+
+
+LINE_STOP_FLAG = ['.', '!', '?', '。', '！', '？',"：", ":", ")", "）", ";"]
+INLINE_EQUATION = 'inline_equation'
+INTER_EQUATION = "displayed_equation"
+TEXT = "text"
+
+def __add_line_period(blocks, layout_bboxes):
+    """
+    为每行添加句号
+    如果这个行
+    1. 以行内公式结尾，但没有任何标点符号,此时加个句号，认为他就是段落结尾。
+    """
+    for block in blocks:
+        for line in block['lines']:
+            last_span = line['spans'][-1]
+            span_type = last_span['type']
+            if span_type in [TEXT, INLINE_EQUATION]:
+                span_content = last_span['content'].strip()
+                if span_type==INLINE_EQUATION and span_content[-1] not in LINE_STOP_FLAG:
+                    if span_type in [INLINE_EQUATION, INTER_EQUATION]:
+                        last_span['content'] = span_content + '.'
+
+
+
+def __valign_lines(blocks, layout_bboxes):
+    """
+    对齐行的左侧和右侧。
+    扫描行的左侧和右侧，如果x0, x1差距不超过3就强行对齐到所处layout的左右两侧（和layout有一段距离）。
+    3是个经验值，TODO，计算得来
+    
+    """
+    
+    min_distance = 3
+    min_sample = 2
+    
+    for layout_box in layout_bboxes:
+        blocks_in_layoutbox = [b for b in blocks if _is_in(b['bbox'], layout_box['layout_bbox'])]
+        if len(blocks_in_layoutbox)==0:
+            continue
+        
+        x0_lst = np.array([[line['bbox'][0], 0] for block in blocks_in_layoutbox for line in block['lines']])
+        x1_lst = np.array([[line['bbox'][2], 0] for block in blocks_in_layoutbox for line in block['lines']])
+        x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
+        x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
+        x0_uniq_label = np.unique(x0_clusters.labels_)
+        x1_uniq_label = np.unique(x1_clusters.labels_)
+        
+        x0_2_new_val = {} # 存储旧值对应的新值映射
+        x1_2_new_val = {}
+        for label in x0_uniq_label:
+            if label==-1:
+                continue
+            x0_index_of_label = np.where(x0_clusters.labels_==label)
+            x0_raw_val = x0_lst[x0_index_of_label][:,0]
+            x0_new_val = np.min(x0_lst[x0_index_of_label][:,0])
+            x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
+        for label in x1_uniq_label:
+            if label==-1:
+                continue
+            x1_index_of_label = np.where(x1_clusters.labels_==label)
+            x1_raw_val = x1_lst[x1_index_of_label][:,0]
+            x1_new_val = np.max(x1_lst[x1_index_of_label][:,0])
+            x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
+        
+        for block in blocks_in_layoutbox:
+            for line in block['lines']:
+                x0, x1 = line['bbox'][0], line['bbox'][2]
+                if x0 in x0_2_new_val:
+                    line['bbox'][0] = int(x0_2_new_val[x0])
+
+                if x1 in x1_2_new_val:
+                    line['bbox'][2] = int(x1_2_new_val[x1])
+            # 其余对不齐的保持不动
+            
+        # 由于修改了block里的line长度，现在需要重新计算block的bbox
+        for block in blocks_in_layoutbox:
+            block['bbox'] = [min([line['bbox'][0] for line in block['lines']]), 
+                            min([line['bbox'][1] for line in block['lines']]), 
+                            max([line['bbox'][2] for line in block['lines']]), 
+                            max([line['bbox'][3] for line in block['lines']])]
+
+
+def __common_pre_proc(blocks, layout_bboxes):
+    """
+    不分语言的，对文本进行预处理
+    """
+    __add_line_period(blocks, layout_bboxes)
+    __valign_lines(blocks, layout_bboxes)
+    
+
+def __pre_proc_zh_blocks(blocks, layout_bboxes):
+    """
+    对中文文本进行分段预处理
+    """
+    pass
+
+
+def __pre_proc_en_blocks(blocks, layout_bboxes):
+    """
+    对英文文本进行分段预处理
+    """
+    pass
+
+
+def __group_line_by_layout(blocks, layout_bboxes, lang="en"):
+    """
+    每个layout内的行进行聚合
+    """
+    # 因为只是一个block一行目前, 一个block就是一个段落
+    lines_group = []
+    
+    for lyout in layout_bboxes:
+        lines = [line for block in blocks if _is_in(block['bbox'], lyout['layout_bbox']) for line in block['lines']]
+        lines_group.append(lines)
+
+    return lines_group
+    
+
+def __split_para_in_layoutbox(lines_group, layout_bboxes, lang="en", char_avg_len=10):
+    """
+    lines_group 进行行分段——layout内部进行分段。
+    1. 先计算每个group的左右边界。
+    2. 然后根据行末尾特征进行分段。
+        末尾特征：以句号等结束符结尾。并且距离右侧边界有一定距离。
+    
+    """
+    def get_span_text(span):
+        c = span.get('content', '')
+        if len(c)==0:
+            c = span.get('image-path', '')
+            
+        return c
+    
+    paras = []
+    right_tail_distance = 1.5 * char_avg_len
+    for lines in lines_group:
+        if len(lines)==0:
+            continue
+        layout_right = max([line['bbox'][2] for line in lines])
+        para = [] # 元素是line
+        for line in lines:
+            line_text = ''.join([get_span_text(span) for span in line['spans']])
+            #logger.info(line_text)
+            last_span_type = line['spans'][-1]['type']
+            if last_span_type in [TEXT, INLINE_EQUATION]:
+                last_char = line['spans'][-1]['content'][-1]
+                if last_char in LINE_STOP_FLAG or line['bbox'][2] < layout_right - right_tail_distance:
+                    para.append(line)
+                    paras.append(para)
+                    # para_text = ''.join([span['content'] for line in para for span in line['spans']])
+                    # logger.info(para_text)
+                    para = []
+                else: 
+                    para.append(line)
+            else: # 其他，图片、表格、行间公式，各自占一段
+                para.append(line)
+                paras.append(para)
+                # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
+                # logger.info(para_text)
+                para = []
+        if len(para)>0:
+            paras.append(para)
+            # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
+            # logger.info(para_text)
+            para = []
+                    
+    return paras
+            
+
+def __do_split(blocks, layout_bboxes, lang="en"):
+    """
+    根据line和layout情况进行分段
+    先实现一个根据行末尾特征分段的简单方法。
+    """
+    """
+    算法思路：
+    1. 扫描layout里每一行，找出来行尾距离layout有边界有一定距离的行。
+    2. 从上述行中找到末尾是句号等可作为断行标志的行。
+    3. 参照上述行尾特征进行分段。
+    4. 图、表，目前独占一行，不考虑分段。
+    """
+    lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
+    layout_paras = __split_para_in_layoutbox(lines_group, layout_bboxes, lang) # block间连接分段
+    
+    return layout_paras
+    
+    
+def para_split(blocks, layout_bboxes, lang="en"):
+    """
+    根据line和layout情况进行分段
+    """
+    __common_pre_proc(blocks, layout_bboxes)
+    if lang=='en':
+        __do_split(blocks, layout_bboxes, lang)
+    elif lang=='zh':
+        __do_split(blocks, layout_bboxes, lang)
+    
+    splited_blocks = __do_split(blocks, layout_bboxes, lang)
+    
+    return splited_blocks

magic_pdf/pdf_parse_by_model.py

@@ -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(
@@ -282,7 +257,6 @@ def parse_pdf_by_model(
         footnote_bboxes_by_model = parse_footnotes_by_model(page_id, page, model_output_json, md_bookname_save_path, debug_mode=debug_mode)
         # 通过规则识别到的footnote
         footnote_bboxes_by_rule = parse_footnotes_by_rule(remain_text_blocks, page_height, page_id, main_text_font)
-
         """进入pdf过滤器，去掉一些不合理的pdf"""
         is_good_pdf, err = pdf_filter(page, remain_text_blocks, table_bboxes, image_bboxes)
         if not is_good_pdf:

magic_pdf/pdf_parse_by_ocr.py

@@ -1,85 +1,263 @@
 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.para.para_split import para_split
+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, para_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,
+        "para_blocks": para_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合并段落
+        para_blocks = para_split(blocks, layout_bboxes)
+        
+        # 获取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, para_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

magic_pdf/pipeline.py

@@ -3,7 +3,8 @@ import sys
 import time
 from urllib.parse import quote
 
-from magic_pdf.libs.commons import read_file, join_path, parse_bucket_key, formatted_time
+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, s3_image_save_path
 from magic_pdf.libs.drop_reason import DropReason
 from magic_pdf.libs.json_compressor import JsonCompressor
 from magic_pdf.dict2md.mkcontent import mk_nlp_markdown
@@ -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,9 +55,9 @@ 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
+        book_name = f"{data_source}/{file_id}"
 
         # 首页存在超量drawing问题
         # special_pdf_list = ['zlib/zlib_21822650']
@@ -78,9 +101,9 @@ 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
+        book_name = f"{data_source}/{file_id}"
         total_page = pdf_meta["total_page"]
         page_width = pdf_meta["page_width_pts"]
         page_height = pdf_meta["page_height_pts"]
@@ -140,13 +163,13 @@ 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
+        book_name = f"{data_source}/{file_id}"
         title = jso.get('title')
         url_encode_title = quote(title, safe='')
         if data_source != 'scihub':
@@ -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,9 +260,9 @@ 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
+    book_name = f"{data_source}/{file_id}"
 
     # 1.23.22已修复
     # if debug_mode:
@@ -264,7 +287,7 @@ def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
     #     jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
     else:
         try:
-            save_path = "s3://mllm-raw-media/pdf2md_img/"
+            save_path = s3_image_save_path
             image_s3_config = get_s3_config(save_path)
             start_time = time.time()  # 记录开始时间
             # 先打印一下book_name和解析开始的时间
@@ -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 = f"{data_source}/{file_id}"
+    try:
+        save_path = s3_image_save_path
+        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

magic_pdf/pre_proc/detect_footer_by_model.py

@@ -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

magic_pdf/pre_proc/detect_footnote.py

@@ -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

magic_pdf/pre_proc/detect_header.py

@@ -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

magic_pdf/pre_proc/detect_page_number.py

@@ -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

magic_pdf/pre_proc/ocr_cut_image.py

@@ -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

magic_pdf/pre_proc/ocr_detect_layout.py

@@ -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

magic_pdf/pre_proc/ocr_dict_merge.py

@@ -0,0 +1,103 @@
+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):
+    if len(spans) == 0:
+        return []
+    else:
+        # 按照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
+
+
+
+
+
+

magic_pdf/pre_proc/ocr_span_list_modify.py

@@ -0,0 +1,210 @@
+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为空,则不处理
+    if len(spans) == 0:
+        pass
+    else:
+        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

magic_pdf/pre_proc/remove_bbox_overlap.py

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