Merge remote-tracking branch 'origin/master'

feat: add layout

.gitignore

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

demo/demo_test.py

@@ -5,6 +5,7 @@ from pathlib import Path
 
 import click
 
+from magic_pdf.dict2md.mkcontent import mk_mm_markdown
 from magic_pdf.pipeline import (
     meta_scan,
     classify_by_type,
@@ -33,7 +34,7 @@ def get_json_from_local_or_s3(book_name=None):
         s3_config = get_s3_config(json_path)
         file_content = read_file(json_path, s3_config)
         json_str = file_content.decode("utf-8")
-        logger.info(json_str)
+        # logger.info(json_str)
         json_object = json.loads(json_str)
     return json_object
 
@@ -55,14 +56,19 @@ def demo_parse_pdf(book_name=None, start_page_id=0, debug_mode=True):
     write_json_to_local(jso, book_name)
 
     jso_md = pdf_intermediate_dict_to_markdown(jso, debug_mode=debug_mode)
-    md_content = jso_md.get("content")
+    content = jso_md.get("content_list")
+    markdown_content = mk_mm_markdown(content)
     if book_name is not None:
-        save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
-        markdown_save_path = join_path(save_tmp_path, "md", book_name + ".md")
+        save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest", "md", book_name)
+        uni_format_save_path = join_path(save_tmp_path,  "book" + ".json")
+        markdown_save_path = join_path(save_tmp_path,  "book" + ".md")
+        with open(uni_format_save_path, "w", encoding="utf-8") as f:
+            f.write(json.dumps(content, ensure_ascii=False, indent=4))
         with open(markdown_save_path, "w", encoding="utf-8") as f:
-            f.write(md_content)
+            f.write(markdown_content)
+            
     else:
-        logger.info(md_content)
+        logger.info(json.dumps(content, ensure_ascii=False))
 
 
 def demo_save_tables(book_name=None, start_page_id=0, debug_mode=True):

demo/ocr_demo.py

@@ -4,7 +4,9 @@ import os
 from loguru import logger
 from pathlib import Path
 
-from magic_pdf.dict2md.ocr_mkcontent import mk_nlp_markdown, mk_mm_markdown
+from app.common.s3 import get_s3_config
+from demo.demo_test import get_json_from_local_or_s3
+from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_mm_markdown_with_para, ocr_mk_nlp_markdown, ocr_mk_mm_markdown, ocr_mk_mm_standard_format
 from magic_pdf.libs.commons import join_path
 from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
 
@@ -29,40 +31,66 @@ def read_json_file(file_path):
     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.json"
+def ocr_local_parse(ocr_pdf_path, ocr_json_file_path):
     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_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)
+        ocr_parse_core(book_name, ocr_pdf_path, ocr_pdf_model_info)
     except Exception as e:
         logger.exception(e)
+
+
+def ocr_online_parse(book_name, start_page_id=0, debug_mode=True):
+    try:
+        json_object = get_json_from_local_or_s3(book_name)
+        # logger.info(json_object)
+        s3_pdf_path = json_object["file_location"]
+        s3_config = get_s3_config(s3_pdf_path)
+        ocr_pdf_model_info = json_object["doc_layout_result"]
+        ocr_parse_core(book_name, s3_pdf_path, ocr_pdf_model_info, s3_config=s3_config)
+    except Exception as e:
+        logger.exception(e)
+
+
+def ocr_parse_core(book_name, ocr_pdf_path, ocr_pdf_model_info, start_page_id=0, s3_config=None):
+    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,
+        s3_config,
+        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 = ocr_mk_mm_markdown_with_para(pdf_info_dict)
+
+    with open(text_content_save_path, "w", encoding="utf-8") as f:
+        f.write(markdown_content)
+
+    standard_format = ocr_mk_mm_standard_format(pdf_info_dict)
+    standard_format_save_path = f"{save_path_with_bookname}/standard_format.txt"
+    with open(standard_format_save_path, "w", encoding="utf-8") as f:
+        f.write(str(standard_format))
+
+
+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"
+    # ocr_pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.pdf"
+    # ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.json"
+    #ocr_local_parse(ocr_pdf_path, ocr_json_file_path)
+    ocr_online_parse(book_name="美国加州中学教材/edu_00000060")

magic_pdf/dict2md/mkcontent.py

@@ -2,9 +2,15 @@ import math
 from loguru import logger
 
 from magic_pdf.libs.boxbase import find_bottom_nearest_text_bbox, find_top_nearest_text_bbox
+from magic_pdf.libs.ocr_content_type import ContentType
+
+TYPE_INLINE_EQUATION = ContentType.InlineEquation
+TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation
+UNI_FORMAT_TEXT_TYPE = ['text', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']
 
 
-def mk_nlp_markdown(para_dict: dict):
+@DeprecationWarning
+def mk_nlp_markdown_1(para_dict: dict):
     """
     对排序后的bboxes拼接内容
     """
@@ -69,14 +75,14 @@ def __insert_before(content, image_content, target):
     return content
 
 
-
-def mk_mm_markdown(para_dict: dict):
+@DeprecationWarning
+def mk_mm_markdown_1(para_dict: dict):
     """拼装多模态markdown"""
     content_lst = []
     for _, page_info in para_dict.items():
         page_lst = [] # 一个page内的段落列表
         para_blocks = page_info.get("para_blocks")
-        pymu_raw_blocks = page_info.get("preproc_blocks")  
+        pymu_raw_blocks = page_info.get("preproc_blocks")
         
         all_page_images = []
         all_page_images.extend(page_info.get("images",[]))
@@ -137,7 +143,7 @@ def mk_mm_markdown(para_dict: dict):
                                 else:
                                     page_md = __insert_before(page_md, img_content, line_txt)
                             else:
-                                logger.error(f"Can't find the location of image {img['image_path']} in the markdown file")
+                                logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #1")
                 else:# 应当在两个block之间
                     # 找到上方最近的block，如果上方没有就找大下方最近的block
                     top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
@@ -150,7 +156,7 @@ def mk_mm_markdown(para_dict: dict):
                             line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
                             page_md = __insert_before(page_md, img_content, line_txt)
                         else:
-                            logger.error(f"Can't find the location of image {img['image_path']} in the markdown file")
+                            logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #2")
                     
         content_lst.append(page_md)
                     
@@ -158,92 +164,190 @@ def mk_mm_markdown(para_dict: dict):
     content_text = "\n\n".join(content_lst)
 
     return content_text
-    
-    
-@DeprecationWarning
-def mk_mm_markdown_1(para_dict: dict):
+
+
+def __insert_after_para(text, image_path, content_list):
     """
-    得到images和tables变量
+    在content_list中找到text，将image_path作为一个新的node插入到text后面
     """
-    image_all_list = []
+    for i, c in enumerate(content_list):
+        content_type = c.get("type")
+        if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
+            img_node = {
+                "type": "image",
+                "img_path": image_path,
+                "img_alt":"",
+                "img_title":"",
+                "img_caption":""
+            }
+            content_list.insert(i+1, img_node)
+            break
+    else:
+        logger.error(f"Can't find the location of image {image_path} in the markdown file, search target is {text}")
     
+
+
+def __insert_before_para(text, image_path, content_list):
+    """
+    在content_list中找到text，将image_path作为一个新的node插入到text前面
+    """
+    for i, c in enumerate(content_list):
+        content_type = c.get("type")
+        if content_type in  UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
+            img_node = {
+                "type": "image",
+                "img_path": image_path,
+                "img_alt":"",
+                "img_title":"",
+                "img_caption":""
+            }
+            content_list.insert(i, img_node)
+            break
+    else:
+        logger.error(f"Can't find the location of image {image_path} in the markdown file, search target is {text}")
+         
+
+def mk_universal_format(para_dict: dict):
+    """
+    构造统一格式 https://aicarrier.feishu.cn/wiki/FqmMwcH69iIdCWkkyjvcDwNUnTY
+    """
+    content_lst = []
     for _, page_info in para_dict.items():
-        images = page_info.get("images",[])
-        tables = page_info.get("tables",[])
-        image_backup = page_info.get("image_backup", [])  
-        table_backup = page_info.get("table_backup",[]) 
-        all_page_images = []
-        all_page_images.extend(images)
-        all_page_images.extend(image_backup)
-        all_page_images.extend(tables)
-        all_page_images.extend(table_backup)
+        page_lst = [] # 一个page内的段落列表
+        para_blocks = page_info.get("para_blocks")
+        pymu_raw_blocks = page_info.get("preproc_blocks")
         
-        pymu_raw_blocks = page_info.get("pymu_raw_blocks")  
-
-        # 提取每个图片所在位置
-        for image_info in all_page_images:
-            x0_image, y0_image, x1_image, y1_image = image_info['bbox'][:4]
-            image_path = image_info['image_path']
-            
-            # 判断图片处于原始PDF中哪个模块之间
-            image_internal_dict = {}
-            image_external_dict = {}
-            between_dict = {}
+        all_page_images = []
+        all_page_images.extend(page_info.get("images",[]))
+        all_page_images.extend(page_info.get("image_backup", []) )
+        all_page_images.extend(page_info.get("tables",[]))
+        all_page_images.extend(page_info.get("table_backup",[]) )
+        
+        if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
+            for img in all_page_images:
+                content_node = {
+                    "type": "image",
+                    "img_path": img['image_path'],
+                    "img_alt":"",
+                    "img_title":"",
+                    "img_caption":""
+                }
+                page_lst.append(content_node) # TODO 图片顺序
+        else:
+            for block in para_blocks:
+                item = block["paras"]
+                for _, p in item.items():
+                    font_type = p['para_font_type']# 对于文本来说，要么是普通文本，要么是个行间公式
+                    if font_type == TYPE_INTERLINE_EQUATION:
+                        content_node = {
+                            "type": "equation",
+                            "latex": p["para_text"]
+                        }
+                        page_lst.append(content_node)
+                    else:
+                        para_text = p["para_text"]
+                        is_title = p["is_para_title"]
+                        title_level = p['para_title_level']
+                        
+                        if is_title:
+                            content_node = {
+                                "type": f"h{title_level}",
+                                "text": para_text
+                            }
+                            page_lst.append(content_node)
+                        else:
+                            content_node = {
+                                "type": "text",
+                                "text": para_text
+                            }
+                            page_lst.append(content_node)
+                            
+        content_lst.extend(page_lst)
+        
+        """插入图片"""
+        for img in all_page_images:
+            imgbox = img['bbox']
+            img_content = f"{img['image_path']}"
+            # 先看在哪个block内
             for block in pymu_raw_blocks:
-                x0, y0, x1, y1 = block['bbox'][:4]
-
-                # 在某个模块内部
-                if x0 <= x0_image < x1 and y0 <= y0_image < y1:
-                    image_internal_dict['bbox'] = [x0_image, y0_image, x1_image, y1_image]
-                    image_internal_dict['path'] = image_path
-                    
-                    # 确定图片在哪句文本之前
-                    y_pre = 0
-                    for line in block['lines']:
-                        x0, y0, x1, y1 = line['spans'][0]['bbox']
-                        if x0 <= x0_image < x1 and y_pre <= y0_image < y0: 
-                            text = line['spans']['text']
-                            image_internal_dict['text'] = text
-                            image_internal_dict['markdown_image'] = f'![image_path]({image_path})'
+                bbox = block['bbox']
+                if bbox[0]-1 <= imgbox[0] < bbox[2]+1 and bbox[1]-1 <= imgbox[1] < bbox[3]+1:# 确定在这个大的block内，然后进入逐行比较距离
+                    for l in block['lines']:
+                        line_box = l['bbox']
+                        if line_box[0]-1 <= imgbox[0] < line_box[2]+1 and line_box[1]-1 <= imgbox[1] < line_box[3]+1: # 在line内的，插入line前面
+                            line_txt = "".join([s['text'] for s in l['spans']])
+                            __insert_before_para(line_txt, img_content, content_lst)
+                            break
+                        break
+                    else:# 在行与行之间
+                        # 找到图片x0,y0与line的x0,y0最近的line
+                        min_distance = 100000
+                        min_line = None
+                        for l in block['lines']:
+                            line_box = l['bbox']
+                            distance = math.sqrt((line_box[0] - imgbox[0])**2 + (line_box[1] - imgbox[1])**2)
+                            if distance < min_distance:
+                                min_distance = distance
+                                min_line = l
+                        if min_line:
+                            line_txt = "".join([s['text'] for s in min_line['spans']])
+                            img_h = imgbox[3] - imgbox[1]
+                            if min_distance<img_h: # 文字在图片前面
+                                __insert_after_para(line_txt, img_content, content_lst)
+                            else:
+                                __insert_before_para(line_txt, img_content, content_lst) 
                             break
                         else:
-                            y_pre = y0
-                # 在某两个模块之间
-                elif x0 <= x0_image < x1:
-                    distance = math.sqrt((x1_image - x0)**2 + (y1_image - y0)**2)
-                    between_dict[block['number']] = distance
-            
-            # 找到与定位点距离最小的文本block
-            if between_dict:
-                min_key = min(between_dict, key=between_dict.get)
-                spans_list = []
-                for span in pymu_raw_blocks[min_key]['lines']: 
-                    for text_piece in span['spans']:
-                        # 防止索引定位文本内容过多
-                        if len(spans_list) < 60:
-                            spans_list.append(text_piece['text'])
-                text1 = ''.join(spans_list)
-                
-                image_external_dict['bbox'] = [x0_image, y0_image, x1_image, y1_image]
-                image_external_dict['path'] = image_path 
-                image_external_dict['text'] = text1
-                image_external_dict['markdown_image'] = f'![image_path]({image_path})'
+                            logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #1")
+            else:# 应当在两个block之间
+                # 找到上方最近的block，如果上方没有就找大下方最近的block
+                top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
+                if top_txt_block:
+                    line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
+                    __insert_after_para(line_txt, img_content, content_lst) 
+                else:
+                    bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, imgbox)
+                    if bottom_txt_block:
+                        line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
+                        __insert_before_para(line_txt, img_content, content_lst) 
+                    else: # TODO ，图片可能独占一列，这种情况上下是没有图片的
+                        logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #2")
+    # end for
+    return content_lst
 
-            # 将内部图片或外部图片存入当页所有图片的列表
-            if len(image_internal_dict) != 0:
-                image_all_list.append(image_internal_dict)
-            elif len(image_external_dict) != 0:
-                image_all_list.append(image_external_dict)
+
+def mk_mm_markdown(content_list):
+    """
+    基于同一格式的内容列表，构造markdown，含图片
+    """
+    content_md = []
+    for c in content_list:
+        content_type = c.get("type")
+        if content_type == "text":
+            content_md.append(c.get("text"))
+        elif content_type == "equation":
+            content = c.get("latex")
+            if content.startswith("$$") and content.endswith("$$"):
+                content_md.append(content)
             else:
-                logger.error(f"Can't find the location of image {image_path} in the markdown file")
+                content_md.append(f"\n$$\n{c.get('latex')}\n$$\n")
+        elif content_type in UNI_FORMAT_TEXT_TYPE:
+            content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
+        elif content_type == "image":
+            content_md.append(f"![]({c.get('img_path')})")
+    return "\n\n".join(content_md)
 
-    content_text = mk_nlp_markdown(para_dict)
-
-    for image_info_extract in image_all_list:
-        loc = __find_index(content_text, image_info_extract['text'])
-        if loc is not None:
-            content_text = __insert_string(content_text, image_info_extract['markdown_image'], loc)
-        else:
-            logger.error(f"Can't find the location of image {image_info_extract['path']} in the markdown file")
-
-    return content_text
+def mk_nlp_markdown(content_list):
+    """
+    基于同一格式的内容列表，构造markdown，不含图片
+    """
+    content_md = []
+    for c in content_list:
+        content_type = c.get("type")
+        if content_type == "text":
+            content_md.append(c.get("text"))
+        elif content_type == "equation":
+            content_md.append(f"$$\n{c.get('latex')}\n$$")
+        elif content_type in UNI_FORMAT_TEXT_TYPE:
+            content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
+    return "\n\n".join(content_md)

magic_pdf/dict2md/ocr_mkcontent.py

@@ -1,4 +1,19 @@
+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
+import wordninja
+import re
+
+
+def split_long_words(text):
+    segments = text.split(' ')
+    for i in range(len(segments)):
+        words = re.findall(r'\w+|[^\w\s]', segments[i], re.UNICODE)
+        for j in range(len(words)):
+            if len(words[j]) > 15:
+                words[j] = ' '.join(wordninja.split(words[j]))
+        segments[i] = ''.join(words)
+    return ' '.join(segments)
 
 
 def ocr_mk_nlp_markdown(pdf_info_dict: dict):
@@ -14,7 +29,7 @@ def ocr_mk_nlp_markdown(pdf_info_dict: dict):
                 for span in line['spans']:
                     if not span.get('content'):
                         continue
-                    content = span['content'].replace('$', '\$')  # 转义$
+                    content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                     if span['type'] == ContentType.InlineEquation:
                         content = f"${content}$"
                     elif span['type'] == ContentType.InterlineEquation:
@@ -26,7 +41,6 @@ def ocr_mk_nlp_markdown(pdf_info_dict: dict):
 
 
 def ocr_mk_mm_markdown(pdf_info_dict: dict):
-
     markdown = []
 
     for _, page_info in pdf_info_dict.items():
@@ -41,9 +55,9 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
                         if not span.get('image_path'):
                             continue
                         else:
-                            content = f"![]({span['image_path']})"
+                            content = f"![]({join_path(s3_image_save_path, span['image_path'])})"
                     else:
-                        content = span['content'].replace('$', '\$')  # 转义$
+                        content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                         if span['type'] == ContentType.InlineEquation:
                             content = f"${content}$"
                         elif span['type'] == ContentType.InterlineEquation:
@@ -52,3 +66,106 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
                 # 在行末添加两个空格以强制换行
                 markdown.append(line_text.strip() + '  ')
     return '\n'.join(markdown)
+
+
+def ocr_mk_mm_markdown_with_para(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:
+                        content = split_long_words(span['content'])
+                        # content = span['content']
+                    elif span_type == ContentType.InlineEquation:
+                        content = f"${span['content']}$"
+                    elif span_type == ContentType.InterlineEquation:
+                        content = f"\n$$\n{span['content']}\n$$\n"
+                    elif span_type in [ContentType.Image, ContentType.Table]:
+                        content = f"\n![]({join_path(s3_image_save_path, span['image_path'])})\n"
+                    para_text += content + ' '
+            markdown.append(para_text.strip() + '  ')
+
+    return '\n\n'.join(markdown)
+
+
+def make_standard_format_with_para(pdf_info_dict: dict):
+    content_list = []
+    for _, page_info in pdf_info_dict.items():
+        paras = page_info.get("para_blocks")
+        if not paras:
+            continue
+        for para in paras:
+            for line in para:
+                content = line_to_standard_format(line)
+                content_list.append(content)
+    return content_list
+
+
+def line_to_standard_format(line):
+    line_text = ""
+    inline_equation_num = 0
+    for span in line['spans']:
+        if not span.get('content'):
+            if not span.get('image_path'):
+                continue
+            else:
+                if span['type'] == ContentType.Image:
+                    content = {
+                        'type': 'image',
+                        'img_path': join_path(s3_image_save_path, span['image_path'])
+                    }
+                    return content
+                elif span['type'] == ContentType.Table:
+                    content = {
+                        'type': 'table',
+                        'img_path': join_path(s3_image_save_path, span['image_path'])
+                    }
+                    return content
+        else:
+            if span['type'] == ContentType.InterlineEquation:
+                interline_equation = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
+                content = {
+                    'type': 'equation',
+                    'latex': f"$$\n{interline_equation}\n$$"
+                }
+                return content
+            elif span['type'] == ContentType.InlineEquation:
+                inline_equation = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
+                line_text += f"${inline_equation}$"
+                inline_equation_num += 1
+            elif span['type'] == ContentType.Text:
+                text_content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
+                line_text += text_content
+    content = {
+        'type': 'text',
+        'text': line_text,
+        'inline_equation_num': inline_equation_num
+    }
+    return content
+
+
+def ocr_mk_mm_standard_format(pdf_info_dict: dict):
+    '''
+    content_list
+    type         string      image/text/table/equation(行间的单独拿出来，行内的和text合并)
+    latex        string      latex文本字段。
+    text         string      纯文本格式的文本数据。
+    md           string      markdown格式的文本数据。
+    img_path     string      s3://full/path/to/img.jpg
+    '''
+    content_list = []
+    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']:
+                content = line_to_standard_format(line)
+                content_list.append(content)
+    return content_list

magic_pdf/libs/commons.py

@@ -16,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):
     """

magic_pdf/libs/draw_bbox.py

@@ -27,7 +27,7 @@ def draw_bbox_with_number(i, bbox_list, page, rgb_config):
         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):
+def draw_layout_bbox(pdf_info_dict, pdf_bytes, out_path):
     layout_bbox_list = []
     dropped_bbox_list = []
     for page in pdf_info_dict.values():
@@ -40,15 +40,14 @@ def draw_layout_bbox(pdf_info_dict, input_path, out_path):
             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):
+    pdf_docs = fitz.open("pdf", pdf_bytes)
+    for i, page in enumerate(pdf_docs):
         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")
+    pdf_docs.save(f"{out_path}/layout.pdf")
 
-def draw_text_bbox(pdf_info_dict, input_path, out_path):
+def draw_text_bbox(pdf_info_dict, pdf_bytes, out_path):
     text_list = []
     inline_equation_list = []
     interline_equation_list = []
@@ -68,13 +67,12 @@ def draw_text_bbox(pdf_info_dict, input_path, out_path):
         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):
+    pdf_docs = fitz.open("pdf", pdf_bytes)
+    for i, page in enumerate(pdf_docs):
         # 获取当前页面的数据
         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")
+    pdf_docs.save(f"{out_path}/text.pdf")

magic_pdf/libs/drop_tag.py

@@ -1,2 +1,18 @@
 
-COLOR_BG_HEADER_TXT_BLOCK = "color_background_header_txt_block"
+COLOR_BG_HEADER_TXT_BLOCK = "color_background_header_txt_block"
+PAGE_NO = "page-no" # 页码
+CONTENT_IN_FOOT_OR_HEADER = 'in-foot-header-area' # 页眉页脚内的文本
+VERTICAL_TEXT = 'vertical-text' # 垂直文本
+ROTATE_TEXT = 'rotate-text' # 旋转文本
+EMPTY_SIDE_BLOCK = 'empty-side-block' # 边缘上的空白没有任何内容的block
+ON_IMAGE_TEXT = 'on-image-text' # 文本在图片上
+ON_TABLE_TEXT = 'on-table-text' # 文本在表格上
+
+
+class DropTag:
+    PAGE_NUMBER = "page_no"
+    HEADER = "header"
+    FOOTER = "footer"
+    FOOTNOTE = "footnote"
+    NOT_IN_LAYOUT = "not_in_layout"
+    SPAN_OVERLAP = "span_overlap"

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/para/para_split.py

@@ -0,0 +1,476 @@
+from sklearn.cluster import DBSCAN
+import numpy as np
+from loguru import logger
+
+from magic_pdf.libs.boxbase import _is_in_or_part_overlap
+from magic_pdf.libs.ocr_content_type import ContentType
+
+
+LINE_STOP_FLAG = ['.', '!', '?', '。', '！', '？',"：", ":", ")", "）", ";"]
+INLINE_EQUATION = ContentType.InlineEquation
+INTERLINE_EQUATION = ContentType.InterlineEquation
+TEXT = "text"
+
+
+def __get_span_text(span):
+    c = span.get('content', '')
+    if len(c)==0:
+        c = span.get('image_path', '')
+        
+    return c
+    
+    
+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 [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, INTERLINE_EQUATION]:
+                        last_span['content'] = span_content + '.'
+
+
+def __detect_line_align_direction(line, new_layout_bboxes):
+    """
+    探测line是左对齐，还是右对齐，还是居中。
+    """
+    lbox = line['bbox']
+    x0, x1 = lbox[0], lbox[2]
+    layout_x0, layout_x1 = new_layout_bboxes[0], new_layout_bboxes[2]
+    if x0 <= layout_x0 and x1 < layout_x1:
+        return "left"
+    elif x0 > layout_x0 and x1 >= layout_x1:
+        return "right"
+    else:
+        return "center"
+    
+def __detect_line_group_align_direction(lines, new_layout_bboxes):
+    """
+    首先把lines按照行距离分成几部分。针对每一部分分别探测。
+    最后返回[(dir, lines), (dir, lines), ...]
+    """
+    pass
+
+def __detect_list_lines(lines, new_layout_bboxes, lang='en'):
+    """
+    探测是否包含了列表，并且把列表的行分开.
+    这样的段落特点是，顶格字母大写/数字，紧跟着几行缩进的。缩进的行首字母含小写的。
+    """
+    def find_repeating_patterns(lst):
+        indices = []
+        ones_indices = []
+        i = 0
+        while i < len(lst) - 1:  # 确保余下元素至少有2个
+            if lst[i] == 1 and lst[i+1] in [2, 3]:  # 额外检查以防止连续出现的1
+                start = i
+                ones_in_this_interval = [i]
+                i += 1
+                while i < len(lst) and lst[i] in [2, 3]:
+                    i += 1
+                # 验证下一个序列是否符合条件
+                if i < len(lst) - 1 and lst[i] == 1 and lst[i+1] in [2, 3] and lst[i-1] in [2, 3]:
+                    while i < len(lst) and lst[i] in [1, 2, 3]:
+                        if lst[i] == 1:
+                            ones_in_this_interval.append(i)
+                        i += 1
+                    indices.append((start, i - 1))
+                    ones_indices.append(ones_in_this_interval)
+                else:
+                    i += 1
+            else:
+                i += 1
+        return indices, ones_indices
+    """===================="""
+    def split_indices(slen, index_array):
+        result = []
+        last_end = 0
+        
+        for start, end in sorted(index_array):
+            if start > last_end:
+                # 前一个区间结束到下一个区间开始之间的部分标记为"text"
+                result.append(('text', last_end, start - 1))
+            # 区间内标记为"list"
+            result.append(('list', start, end))
+            last_end = end + 1
+
+        if last_end < slen:
+            # 如果最后一个区间结束后还有剩余的字符串，将其标记为"text"
+            result.append(('text', last_end, slen - 1))
+
+        return result
+    """===================="""
+
+    if lang!='en':
+        return lines, None
+    else:
+        total_lines = len(lines)
+        line_fea_encode = []
+        """
+        对每一行进行特征编码，编码规则如下：
+        1. 如果行顶格，且大写字母开头或者数字开头，编码为1
+        2. 如果顶格，其他非大写开头编码为4
+        3. 如果非顶格，首字符大写，编码为2
+        4. 如果非顶格，首字符非大写编码为3
+        """
+        for l in lines:
+            first_char = __get_span_text(l['spans'][0])[0]
+            layout_left = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)[0]
+            if l['bbox'][0] == layout_left:
+                if first_char.isupper() or first_char.isdigit():
+                    line_fea_encode.append(1)
+                else:
+                    line_fea_encode.append(4)
+            else:
+                if first_char.isupper():
+                    line_fea_encode.append(2)
+                else:
+                    line_fea_encode.append(3)
+                    
+        # 然后根据编码进行分段, 选出来 1,2,3连续出现至少2次的行，认为是列表。
+        
+        list_indice, list_start_idx  = find_repeating_patterns(line_fea_encode)
+        if len(list_indice)>0:
+            logger.info(f"发现了列表，列表行数：{list_indice}， {list_start_idx}")
+        
+        # TODO check一下这个特列表里缩进的行左侧是不是对齐的。
+        segments = []
+        for start, end in list_indice:
+            for i in range(start, end+1):
+                if i>0:
+                    if line_fea_encode[i] == 4:
+                        logger.info(f"列表行的第{i}行不是顶格的")
+                        break
+            else:
+                logger.info(f"列表行的第{start}到第{end}行是列表")
+        
+        return split_indices(total_lines, list_indice), list_start_idx
+        
+            
+
+def __valign_lines(blocks, layout_bboxes):
+    """
+    在一个layoutbox内对齐行的左侧和右侧。
+    扫描行的左侧和右侧，如果x0, x1差距不超过一个阈值，就强行对齐到所处layout的左右两侧（和layout有一段距离）。
+    3是个经验值，TODO，计算得来，可以设置为1.5个正文字符。
+    """
+    
+    min_distance = 3
+    min_sample = 2
+    new_layout_bboxes = []
+    
+    for layout_box in layout_bboxes:
+        blocks_in_layoutbox = [b for b in blocks if _is_in_or_part_overlap(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']])]
+            
+        """新计算layout的bbox，因为block的bbox变了。"""
+        layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
+        layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
+        layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
+        layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
+        new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])
+            
+    return new_layout_bboxes
+
+
+def __common_pre_proc(blocks, layout_bboxes):
+    """
+    不分语言的，对文本进行预处理
+    """
+    #__add_line_period(blocks, layout_bboxes)
+    aligned_layout_bboxes = __valign_lines(blocks, layout_bboxes)
+    
+    return aligned_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_or_part_overlap(block['bbox'], lyout['layout_bbox']) for line in block['lines']]
+        lines_group.append(lines)
+
+    return lines_group
+    
+
+def __split_para_in_layoutbox(lines_group, new_layout_bbox, lang="en", char_avg_len=10):
+    """
+    lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
+    1. 先计算每个group的左右边界。
+    2. 然后根据行末尾特征进行分段。
+        末尾特征：以句号等结束符结尾。并且距离右侧边界有一定距离。
+        且下一行开头不留空白。
+    
+    """
+    line_group_end_with_list = [] # 这个layout最后是不是列表，用于跨layout列表合并
+    paras = []
+    right_tail_distance = 1.5 * char_avg_len
+    for lines in lines_group:
+        total_lines = len(lines)
+        if total_lines<=1: # 0行无需处理。1行无法分段。
+            continue
+        
+        """在进入到真正的分段之前，要对文字块从统计维度进行对齐方式的探测，
+            对齐方式分为以下：
+            1. 左对齐的文本块(特点是左侧顶格，或者左侧不顶格但是右侧顶格的行数大于非顶格的行数，顶格的首字母有大写也有小写)
+                1) 右侧对齐的行，单独成一段
+                2) 中间对齐的行，按照字体/行高聚合成一段
+            2. 左对齐的列表块（其特点是左侧顶格的行数小于等于非顶格的行数，非定格首字母会有小写，顶格90%是大写。并且左侧顶格行数大于1，大于1是为了这种模式连续出现才能称之为列表）
+                这样的文本块，顶格的为一个段落开头，紧随其后非顶格的行属于这个段落。
+        """
+        
+        text_segments, list_start_line = __detect_list_lines(lines, new_layout_bbox, lang)
+        """根据list_range，把lines分成几个部分
+        
+        """
+        
+        layout_right = __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[2]
+        layout_left = __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[0]
+        para = [] # 元素是line
+        is_lines_end_with_list = False
+        for content_type, start, end in text_segments:
+            if content_type == 'list':
+                for i, line in enumerate(lines[start:end+1]):
+                    line_x0 = line['bbox'][0]
+                    if line_x0 == layout_left: # 列表开头
+                        if len(para)>0:
+                            paras.append(para)
+                            para = []
+                        para.append(line)
+                    else:
+                        para.append(line)
+                if len(para)>0:
+                    paras.append(para)
+                    para = []
+                is_lines_end_with_list = True
+            else:
+                for i, line in enumerate(lines[start:end+1]):
+                    # 如果i有下一行，那么就要根据下一行位置综合判断是否要分段。如果i之后没有行，那么只需要判断一下行结尾特征。
+                    cur_line_type = line['spans'][-1]['type']
+                    next_line = lines[i+1] if i<total_lines-1 else None
+                    
+                    if cur_line_type in [TEXT, INLINE_EQUATION]:
+                        if line['bbox'][2] < layout_right - right_tail_distance:
+                            para.append(line)
+                            paras.append(para)
+                            para = []
+                        elif line['bbox'][2] >= layout_right - right_tail_distance and next_line and next_line['bbox'][0] == layout_left: # 现在这行到了行尾沾满，下一行存在且顶格。
+                            para.append(line)
+                        else: 
+                            para.append(line)
+                            paras.append(para)
+                            para = []
+                    else: # 其他，图片、表格、行间公式，各自占一段
+                        if len(para)>0:  # 先把之前的段落加入到结果中
+                            paras.append(para)
+                            para = []
+                        paras.append([line]) # 再把当前行加入到结果中。当前行为行间公式、图、表等。
+                        para = []
+                        
+                if len(para)>0:
+                    paras.append(para)
+                    para = []
+                is_lines_end_with_list = False
+                
+        line_group_end_with_list.append(is_lines_end_with_list)
+                
+                    
+    return paras, line_group_end_with_list
+
+
+def __find_layout_bbox_by_line(line_bbox, layout_bboxes):
+    """
+    根据line找到所在的layout
+    """
+    for layout in layout_bboxes:
+        if _is_in_or_part_overlap(line_bbox, layout):
+            return layout
+    return None
+
+
+def __connect_para_inter_layoutbox(layout_paras, new_layout_bbox, line_group_end_with_list, lang="en"):
+    """
+    layout之间进行分段。
+    主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
+    连接的条件需要同时满足：
+    1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
+    2. 下一行开头不留空白。
+
+    """
+    connected_layout_paras = []
+    for i, para in enumerate(layout_paras):
+        if i==0:
+            connected_layout_paras.append(para)
+            continue
+        pre_last_line = layout_paras[i-1][-1]
+        next_first_line = layout_paras[i][0]
+        pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
+        pre_last_line_type = pre_last_line['spans'][-1]['type']
+        next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
+        next_first_line_type = next_first_line['spans'][0]['type']
+        if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]: # TODO，真的要做好，要考虑跨table, image, 行间的情况
+            connected_layout_paras.append(para)
+            continue
+        
+        
+        pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)[2]
+        next_x0_min = __find_layout_bbox_by_line(next_first_line['bbox'], new_layout_bbox)[0]
+        
+        pre_last_line_text = pre_last_line_text.strip()
+        next_first_line_text = next_first_line_text.strip()
+        if pre_last_line['bbox'][2] == pre_x2_max and pre_last_line_text[-1] not in LINE_STOP_FLAG and next_first_line['bbox'][0]==next_x0_min: # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
+            """连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
+            connected_layout_paras[-1].extend(para)
+        else:
+            """连接段落条件不成立，将前一个layout的段落加入到结果中。"""
+            connected_layout_paras.append(para)
+    
+    return connected_layout_paras
+
+
+def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, lang):
+    """
+    连接起来相邻两个页面的段落——前一个页面最后一个段落和后一个页面的第一个段落。
+    是否可以连接的条件：
+    1. 前一个页面的最后一个段落最后一行沾满整个行。并且没有结尾符号。
+    2. 后一个页面的第一个段落第一行没有空白开头。
+    """
+    # 有的页面可能压根没有文字
+    if len(pre_page_paras)==0 or len(next_page_paras)==0:
+        return False
+    pre_last_para = pre_page_paras[-1]
+    next_first_para = next_page_paras[0]
+    pre_last_line = pre_last_para[-1]
+    next_first_line = next_first_para[0]
+    pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
+    pre_last_line_type = pre_last_line['spans'][-1]['type']
+    next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
+    next_first_line_type = next_first_line['spans'][0]['type']
+    
+    if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]: # TODO，真的要做好，要考虑跨table, image, 行间的情况
+        # 不是文本，不连接
+        return False
+    
+    pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], pre_page_layout_bbox)[2]
+    next_x0_min = __find_layout_bbox_by_line(next_first_line['bbox'], next_page_layout_bbox)[0]
+    
+    pre_last_line_text = pre_last_line_text.strip()
+    next_first_line_text = next_first_line_text.strip()
+    if pre_last_line['bbox'][2] == pre_x2_max and pre_last_line_text[-1] not in LINE_STOP_FLAG and next_first_line['bbox'][0]==next_x0_min: # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
+        """连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
+        pre_page_paras[-1].extend(next_first_para)
+        next_page_paras.pop(0) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
+        return True
+    else:
+        return False
+
+
+def __do_split(blocks, layout_bboxes, new_layout_bbox, lang="en"):
+    """
+    根据line和layout情况进行分段
+    先实现一个根据行末尾特征分段的简单方法。
+    """
+    """
+    算法思路：
+    1. 扫描layout里每一行，找出来行尾距离layout有边界有一定距离的行。
+    2. 从上述行中找到末尾是句号等可作为断行标志的行。
+    3. 参照上述行尾特征进行分段。
+    4. 图、表，目前独占一行，不考虑分段。
+    """
+    lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
+    layout_paras, line_group_end_with_list = __split_para_in_layoutbox(lines_group, new_layout_bbox, lang) # layout内分段
+    connected_layout_paras = __connect_para_inter_layoutbox(layout_paras, new_layout_bbox, line_group_end_with_list, lang) # layout间链接段落
+    return connected_layout_paras
+    
+    
+def para_split(pdf_info_dict, lang="en"):
+    """
+    根据line和layout情况进行分段
+    """
+    new_layout_of_pages = [] # 数组的数组，每个元素是一个页面的layoutS
+    for _, page in pdf_info_dict.items():
+        blocks = page['preproc_blocks']
+        layout_bboxes = page['layout_bboxes']
+        new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
+        new_layout_of_pages.append(new_layout_bbox)
+        splited_blocks = __do_split(blocks, layout_bboxes, new_layout_bbox, lang)
+        page['para_blocks'] = splited_blocks
+        
+    """连接页面与页面之间的可能合并的段落"""
+    pdf_infos = list(pdf_info_dict.values())
+    for i, page in enumerate(pdf_info_dict.values()):
+        if i==0:
+            continue
+        pre_page_paras = pdf_infos[i-1]['para_blocks']
+        next_page_paras = pdf_infos[i]['para_blocks']
+        pre_page_layout_bbox = new_layout_of_pages[i-1]
+        next_page_layout_bbox = new_layout_of_pages[i]
+        
+        is_conn= __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, lang) 
+        if is_conn:
+            logger.info(f"连接了第{i-1}页和第{i}页的段落")

magic_pdf/pdf_parse_by_model.py

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

@@ -14,8 +14,10 @@ from magic_pdf.libs.commons import (
     get_docx_model_output,
 )
 from magic_pdf.libs.coordinate_transform import get_scale_ratio
+from magic_pdf.libs.drop_tag import DropTag
 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
@@ -33,7 +35,7 @@ from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
 
 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,
+                             dropped_text_block, dropped_image_block, dropped_table_block, dropped_equation_block,
                              need_remove_spans_bboxes_dict):
     return_dict = {
         'preproc_blocks': blocks,
@@ -48,6 +50,7 @@ def construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, lay
         'droped_text_block': dropped_text_block,
         'droped_image_block': dropped_image_block,
         'droped_table_block': dropped_table_block,
+        'dropped_equation_block': dropped_equation_block,
         'droped_bboxes': need_remove_spans_bboxes_dict,
     }
     return return_dict
@@ -131,10 +134,10 @@ def parse_pdf_by_ocr(
 
         # 构建需要remove的bbox字典
         need_remove_spans_bboxes_dict = {
-            "page_no": page_no_bboxes,
-            "header": header_bboxes,
-            "footer": footer_bboxes,
-            "footnote": footnote_bboxes,
+            DropTag.PAGE_NUMBER: page_no_bboxes,
+            DropTag.HEADER: header_bboxes,
+            DropTag.FOOTER: footer_bboxes,
+            DropTag.FOOTNOTE: footnote_bboxes,
         }
 
         layout_dets = ocr_page_info["layout_dets"]
@@ -156,6 +159,9 @@ def parse_pdf_by_ocr(
                     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', # 页码
@@ -193,16 +199,19 @@ def parse_pdf_by_ocr(
             else:
                 continue
 
+
+
+
         # 删除重叠spans中较小的那些
-        spans = remove_overlaps_min_spans(spans)
+        spans, dropped_spans_by_span_overlap = 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)
+        spans, dropped_spans_by_removed_bboxes = 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)
+        spans = cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client)
 
         # 行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)
         displayed_list = []
@@ -222,24 +231,50 @@ def parse_pdf_by_ocr(
         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, dropped_spans_by_layout = 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)
 
+        # drop的span_list合并
+        dropped_spans = []
+        dropped_spans.extend(dropped_spans_by_span_overlap)
+        dropped_spans.extend(dropped_spans_by_removed_bboxes)
+        dropped_spans.extend(dropped_spans_by_layout)
+
+        dropped_text_block = []
+        dropped_image_block = []
+        dropped_table_block = []
+        dropped_equation_block = []
+        for span in dropped_spans:
+            # drop出的spans进行分类
+            if span['type'] == ContentType.Text:
+                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)
+            elif span['type'] in [ContentType.InlineEquation, ContentType.InterlineEquation]:
+                dropped_equation_block.append(span)
+
+
+
         # 构造pdf_info_dict
         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,
+                                             dropped_text_block, dropped_image_block, dropped_table_block, dropped_equation_block,
                                              need_remove_spans_bboxes_dict)
         pdf_info_dict[f"page_{page_id}"] = page_info
 
+    """分段"""
+    para_split(pdf_info_dict)
+    
     # 在测试时,保存调试信息
     if debug_mode:
         params_file_save_path = join_path(
@@ -249,7 +284,7 @@ def parse_pdf_by_ocr(
             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)
+        draw_layout_bbox(pdf_info_dict, pdf_bytes, md_bookname_save_path)
+        draw_text_bbox(pdf_info_dict, pdf_bytes, md_bookname_save_path)
 
     return pdf_info_dict

magic_pdf/pdf_parse_for_train.py

@@ -0,0 +1,684 @@
+import time
+
+# from anyio import Path
+
+from magic_pdf.libs.commons import (
+    fitz,
+    get_delta_time,
+    get_img_s3_client,
+    get_docx_model_output,
+)
+import json
+import os
+from copy import deepcopy
+import math
+from loguru import logger
+from magic_pdf.layout.bbox_sort import (
+    prepare_bboxes_for_layout_split,
+)
+from magic_pdf.layout.layout_sort import (
+    LAYOUT_UNPROC,
+    get_bboxes_layout,
+    get_columns_cnt_of_layout,
+    sort_text_block,
+)
+from magic_pdf.libs.drop_reason import DropReason
+from magic_pdf.libs.markdown_utils import escape_special_markdown_char
+from magic_pdf.libs.safe_filename import sanitize_filename
+from magic_pdf.libs.vis_utils import draw_bbox_on_page, draw_layout_bbox_on_page
+from magic_pdf.pre_proc.detect_images import parse_images
+from magic_pdf.pre_proc.detect_tables import parse_tables  # 获取tables的bbox
+from magic_pdf.pre_proc.detect_equation import parse_equations  # 获取equations的bbox
+from magic_pdf.pre_proc.detect_header import parse_headers  # 获取headers的bbox
+from magic_pdf.pre_proc.detect_page_number import parse_pageNos  # 获取pageNos的bbox
+from magic_pdf.pre_proc.detect_footnote import (
+    parse_footnotes_by_model,
+    parse_footnotes_by_rule,
+)  # 获取footnotes的bbox
+from magic_pdf.pre_proc.detect_footer_by_model import parse_footers  # 获取footers的bbox
+
+from magic_pdf.post_proc.detect_para import (
+    ParaProcessPipeline,
+    TitleDetectionException,
+    TitleLevelException,
+    ParaSplitException,
+    ParaMergeException,
+    DenseSingleLineBlockException,
+)
+from magic_pdf.pre_proc.main_text_font import get_main_text_font
+from magic_pdf.pre_proc.remove_colored_strip_bbox import remove_colored_strip_textblock
+from magic_pdf.pre_proc.remove_footer_header import remove_headder_footer_one_page
+from magic_pdf.train_utils.extract_caption import extract_caption_bbox
+
+"""
+from para.para_pipeline import ParaProcessPipeline
+from para.exceptions import (
+    TitleDetectionException,
+    TitleLevelException,
+    ParaSplitException,
+    ParaMergeException,
+    DenseSingleLineBlockException,
+)
+"""
+
+from magic_pdf.libs.commons import read_file, join_path
+from magic_pdf.libs.pdf_image_tools import save_images_by_bboxes
+from magic_pdf.post_proc.remove_footnote import (
+    merge_footnote_blocks,
+    remove_footnote_blocks,
+)
+from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker
+from magic_pdf.pre_proc.equations_replace import (
+    combine_chars_to_pymudict,
+    remove_chars_in_text_blocks,
+    replace_equations_in_textblock,
+)
+from magic_pdf.pre_proc.pdf_pre_filter import pdf_filter
+from magic_pdf.pre_proc.detect_footer_header_by_statistics import drop_footer_header
+from magic_pdf.pre_proc.construct_paras import construct_page_component
+from magic_pdf.pre_proc.fix_image import (
+    combine_images,
+    fix_image_vertical,
+    fix_seperated_image,
+    include_img_title,
+)
+from magic_pdf.post_proc.pdf_post_filter import pdf_post_filter
+from magic_pdf.pre_proc.remove_rotate_bbox import (
+    get_side_boundry,
+    remove_rotate_side_textblock,
+    remove_side_blank_block,
+)
+from magic_pdf.pre_proc.resolve_bbox_conflict import (
+    check_text_block_horizontal_overlap,
+    resolve_bbox_overlap_conflict,
+)
+from magic_pdf.pre_proc.fix_table import (
+    fix_table_text_block,
+    fix_tables,
+    include_table_title,
+)
+from magic_pdf.pre_proc.solve_line_alien import solve_inline_too_large_interval
+
+denseSingleLineBlockException_msg = DenseSingleLineBlockException().message
+titleDetectionException_msg = TitleDetectionException().message
+titleLevelException_msg = TitleLevelException().message
+paraSplitException_msg = ParaSplitException().message
+paraMergeException_msg = ParaMergeException().message
+
+
+def parse_pdf_for_train(
+    s3_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,
+    junk_img_bojids=[],
+    debug_mode=False,
+):
+    pdf_bytes = read_file(s3_pdf_path, s3_pdf_profile)
+    save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
+    md_bookname_save_path = ""
+    book_name = sanitize_filename(book_name)
+    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 = {}
+    img_s3_client = get_img_s3_client(
+        save_path, image_s3_config
+    )  # 更改函数名和参数，避免歧义
+    # img_s3_client = "img_s3_client"  #不创建这个对象，直接用字符串占位
+
+    start_time = time.time()
+
+    """通过统计pdf全篇文字,识别正文字体"""
+    main_text_font = get_main_text_font(pdf_docs)
+
+    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):
+        page = pdf_docs[page_id]
+        page_width = page.rect.width
+        page_height = 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
+        """
+        # 通过一个规则，过滤掉单页超过1500非junkimg的pdf
+        # 对单页面非重复id的img数量做统计,如果当前页超过1500则直接return need_drop
+        """
+        page_imgs = page.get_images()
+        img_counts = 0
+        for img in page_imgs:
+            img_bojid = img[0]
+            if img_bojid in junk_img_bojids:  # 判断这个图片在不在junklist中
+                continue  # 如果在junklist就不用管了，跳过
+            else:
+                recs = page.get_image_rects(img, transform=True)
+                if recs:  # 如果这张图在当前页面有展示
+                    img_counts += 1
+        if (
+            img_counts >= 1500
+        ):  # 如果去除了junkimg的影响，单页img仍然超过1500的话，就排除当前pdf
+            logger.warning(
+                f"page_id: {page_id}, img_counts: {img_counts}, drop this pdf: {book_name}, drop_reason: {DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS}"
+            )
+            result = {
+                "need_drop": True,
+                "drop_reason": DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS,
+            }
+            if not debug_mode:
+                return result
+
+        """
+        ==================================================================================================================================
+        首先获取基本的block数据，对pdf进行分解，获取图片、表格、公式、text的bbox
+        """
+        # 解析pdf原始文本block
+        text_raw_blocks = page.get_text(
+            "dict",
+            flags=fitz.TEXTFLAGS_TEXT,
+        )["blocks"]
+        model_output_json = get_docx_model_output(
+            pdf_model_output, pdf_model_profile, page_id
+        )
+
+        # 解析图片
+        image_bboxes = parse_images(page_id, page, model_output_json, junk_img_bojids)
+        image_bboxes = fix_image_vertical(
+            image_bboxes, text_raw_blocks
+        )  # 修正图片的位置
+        image_bboxes = fix_seperated_image(image_bboxes)  # 合并有边重合的图片
+
+        old_image_bboxes = deepcopy(image_bboxes)
+        image_bboxes = include_img_title(
+            text_raw_blocks, image_bboxes
+        )  # 向图片上方和下方寻找title，使用规则进行匹配，暂时只支持英文规则
+        """此时image_bboxes中可能出现这种情况，水平并列的2个图片，下方分别有各自的子标题，2个子标题下方又有大标题（形如Figxxx)，会出现2个图片的bbox都包含了这个大标题，这种情况需要把图片合并"""
+        image_bboxes = combine_images(image_bboxes)  # 合并图片
+
+        # 解析表格并对table_bboxes进行位置的微调,防止表格周围的文字被截断
+        table_bboxes = parse_tables(page_id, page, model_output_json)
+        table_bboxes = fix_tables(
+            page, table_bboxes, include_table_title=True, scan_line_num=2
+        )  # 修正
+        table_bboxes = fix_table_text_block(
+            text_raw_blocks, table_bboxes
+        )  # 修正与text block的关系,某些table修正与pymupdf获取到的table内textblock没有完全包含，因此要进行一次修正。
+        # debug_show_bbox(pdf_docs, page_id, table_bboxes, [], [b['bbox'] for b in text_raw_blocks], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 7)
+
+        old_table_bboxes = deepcopy(table_bboxes)
+        table_bboxes = include_table_title(
+            text_raw_blocks, table_bboxes
+        )  # 向table上方和下方寻找title，使用规则进行匹配，暂时只支持英文规则
+
+        # 解析公式
+        equations_inline_bboxes, equations_interline_bboxes = parse_equations(
+            page_id, page, model_output_json
+        )
+
+        # get image box and caption !
+        image_bboxes_with_caption = extract_caption_bbox(image_bboxes, old_image_bboxes)
+
+        # get table box and caption !
+        table_bboxes_with_caption = extract_caption_bbox(table_bboxes, old_table_bboxes)
+
+        """
+        ==================================================================================================================================
+        进入预处理-1阶段
+        -------------------
+        # # 解析标题
+        # title_bboxs = parse_titles(page_id, page, model_output_json)
+        # # 评估Layout是否规整、简单
+        # isSimpleLayout_flag, fullColumn_cnt, subColumn_cnt, curPage_loss = evaluate_pdf_layout(page_id, page, model_output_json)
+        接下来开始进行预处理过程
+        """
+
+        """去掉每页的页码、页眉、页脚"""
+        page_no_bboxs = parse_pageNos(page_id, page, model_output_json)
+        header_bboxs = parse_headers(page_id, page, model_output_json)
+        footer_bboxs = parse_footers(page_id, page, model_output_json)
+        (
+            image_bboxes,
+            table_bboxes,
+            remain_text_blocks,
+            removed_hdr_foot_txt_block,
+            removed_hdr_foot_img_block,
+            removed_hdr_foot_table,
+        ) = remove_headder_footer_one_page(
+            text_raw_blocks,
+            image_bboxes,
+            table_bboxes,
+            header_bboxs,
+            footer_bboxs,
+            page_no_bboxs,
+            page_width,
+            page_height,
+        )
+
+        """去除页面上半部分长条色块内的文本块"""
+        remain_text_blocks, removed_colored_narrow_strip_background_text_block = (
+            remove_colored_strip_textblock(remain_text_blocks, page)
+        )
+
+        # debug_show_bbox(pdf_docs, page_id, footnote_bboxes_by_model, [b['bbox'] for b in remain_text_blocks], header_bboxs, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 7)
+
+        """去掉旋转的文字：水印、垂直排列的文字"""
+        remain_text_blocks, removed_non_horz_text_block = remove_rotate_side_textblock(
+            remain_text_blocks, page_width, page_height
+        )  # 去掉水印，非水平文字
+        remain_text_blocks, removed_empty_side_block = remove_side_blank_block(
+            remain_text_blocks, page_width, page_height
+        )  # 删除页面四周可能会留下的完全空白的textblock，这种block形成原因未知
+
+        """出现在图片、表格上的文字块去掉，把层叠的图片单独分离出来，不参与layout的计算"""
+        (
+            image_bboxes,
+            table_bboxes,
+            equations_interline_bboxes,
+            equations_inline_bboxes,
+            remain_text_blocks,
+            text_block_on_image_removed,
+            images_overlap_backup,
+            interline_eq_temp_text_block,
+        ) = resolve_bbox_overlap_conflict(
+            image_bboxes,
+            table_bboxes,
+            equations_interline_bboxes,
+            equations_inline_bboxes,
+            remain_text_blocks,
+        )
+
+        # """去掉footnote, 从文字和图片中"""
+        # # 通过模型识别到的footnote
+        # 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)
+        """
+        ==================================================================================================================================
+        """
+        if debug_mode:  # debugmode截图到本地
+            save_path = join_path(save_tmp_path, "md")
+
+        # 把图、表、公式都进行截图，保存到存储上，返回图片路径作为内容
+        image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info = (
+            save_images_by_bboxes(
+                book_name,
+                page_id,
+                page,
+                save_path,
+                image_bboxes,
+                images_overlap_backup,
+                table_bboxes,
+                equations_inline_bboxes,
+                equations_interline_bboxes,
+                # 传入img_s3_client
+                img_s3_client,
+            )
+        )  # 只要表格和图片的截图
+
+        """"以下进入到公式替换环节 """
+        char_level_text_blocks = page.get_text("rawdict", flags=fitz.TEXTFLAGS_TEXT)[
+            "blocks"
+        ]
+        remain_text_blocks = combine_chars_to_pymudict(
+            remain_text_blocks, char_level_text_blocks
+        )  # 合并chars
+        remain_text_blocks = replace_equations_in_textblock(
+            remain_text_blocks, inline_eq_info, interline_eq_info
+        )
+        remain_text_blocks = remove_citation_marker(
+            remain_text_blocks
+        )  # 公式替换之后去角标，防止公式无法替换成功。但是这样也会带来个问题就是把角标当公式。各有优劣。
+        remain_text_blocks = remove_chars_in_text_blocks(
+            remain_text_blocks
+        )  # 减少中间态数据体积
+        # debug_show_bbox(pdf_docs, page_id, [b['bbox'] for b in inline_eq_info], [b['bbox'] for b in interline_eq_info], [], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 3)
+
+        """去掉footnote, 从文字和图片中(先去角标再去footnote试试)"""
+        # 通过模型识别到的footnote
+        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:
+            logger.warning(
+                f"page_id: {page_id}, drop this pdf: {book_name}, reason: {err}"
+            )
+            if not debug_mode:
+                return err
+
+        """
+        ==================================================================================================================================
+        进行版面布局切分和过滤
+        """
+        """在切分之前，先检查一下bbox是否有左右重叠的情况，如果有，那么就认为这个pdf暂时没有能力处理好，这种左右重叠的情况大概率是由于pdf里的行间公式、表格没有被正确识别出来造成的 """
+
+        is_text_block_horz_overlap = check_text_block_horizontal_overlap(
+            remain_text_blocks, header_bboxs, footer_bboxs
+        )
+
+        if is_text_block_horz_overlap:
+            # debug_show_bbox(pdf_docs, page_id, [b['bbox'] for b in remain_text_blocks], [], [], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 0)
+            logger.warning(
+                f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TEXT_BLCOK_HOR_OVERLAP}"
+            )
+            result = {
+                "need_drop": True,
+                "drop_reason": DropReason.TEXT_BLCOK_HOR_OVERLAP,
+            }
+            if not debug_mode:
+                return result
+
+        """统一格式化成一个数据结构用于计算layout"""
+        page_y0 = 0 if len(header_bboxs) == 0 else max([b[3] for b in header_bboxs])
+        page_y1 = (
+            page_height if len(footer_bboxs) == 0 else min([b[1] for b in footer_bboxs])
+        )
+        left_x, right_x = get_side_boundry(
+            removed_non_horz_text_block, page_width, page_height
+        )
+        page_boundry = [
+            math.floor(left_x),
+            page_y0 + 1,
+            math.ceil(right_x),
+            page_y1 - 1,
+        ]
+        # 返回的是一个数组，每个元素[x0, y0, x1, y1, block_content, idx_x, idx_y], 初始时候idx_x, idx_y都是None. 对于图片、公式来说，block_content是图片的地址， 对于段落来说，block_content是段落的内容
+
+        all_bboxes = prepare_bboxes_for_layout_split(
+            image_info,
+            image_backup_info,
+            table_info,
+            inline_eq_info,
+            interline_eq_info,
+            remain_text_blocks,
+            page_boundry,
+            page,
+        )
+        # debug_show_bbox(pdf_docs, page_id, [], [], all_bboxes, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 1)
+        """page_y0, page_y1能够过滤掉页眉和页脚，不会算作layout内"""
+        layout_bboxes, layout_tree = get_bboxes_layout(
+            all_bboxes, page_boundry, page_id
+        )
+
+        if (
+            len(remain_text_blocks) > 0
+            and len(all_bboxes) > 0
+            and len(layout_bboxes) == 0
+        ):
+            logger.warning(
+                f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.CAN_NOT_DETECT_PAGE_LAYOUT}"
+            )
+            result = {
+                "need_drop": True,
+                "drop_reason": DropReason.CAN_NOT_DETECT_PAGE_LAYOUT,
+            }
+            if not debug_mode:
+                return result
+
+        """以下去掉复杂的布局和超过2列的布局"""
+        if any(
+            [lay["layout_label"] == LAYOUT_UNPROC for lay in layout_bboxes]
+        ):  # 复杂的布局
+            logger.warning(
+                f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.COMPLICATED_LAYOUT}"
+            )
+            result = {"need_drop": True, "drop_reason": DropReason.COMPLICATED_LAYOUT}
+            if not debug_mode:
+                return result
+
+        layout_column_width = get_columns_cnt_of_layout(layout_tree)
+        if layout_column_width > 2:  # 去掉超过2列的布局pdf
+            logger.warning(
+                f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TOO_MANY_LAYOUT_COLUMNS}"
+            )
+            result = {
+                "need_drop": True,
+                "drop_reason": DropReason.TOO_MANY_LAYOUT_COLUMNS,
+                "extra_info": {"column_cnt": layout_column_width},
+            }
+            if not debug_mode:
+                return result
+
+        """
+        ==================================================================================================================================
+        构造出下游需要的数据结构
+        """
+        remain_text_blocks = (
+            remain_text_blocks + interline_eq_temp_text_block
+        )  # 把计算layout时候临时删除的行间公式再放回去，防止行间公式替换的时候丢失。
+        removed_text_blocks = []
+        removed_text_blocks.extend(removed_hdr_foot_txt_block)
+        # removed_text_blocks.extend(removed_footnote_text_block)
+        removed_text_blocks.extend(text_block_on_image_removed)
+        removed_text_blocks.extend(removed_non_horz_text_block)
+        removed_text_blocks.extend(removed_colored_narrow_strip_background_text_block)
+
+        removed_images = []
+        # removed_images.extend(footnote_imgs)
+        removed_images.extend(removed_hdr_foot_img_block)
+
+        images_backup = []
+        images_backup.extend(image_backup_info)
+        remain_text_blocks = escape_special_markdown_char(
+            remain_text_blocks
+        )  # 转义span里的text
+        sorted_text_remain_text_block = sort_text_block(
+            remain_text_blocks, layout_bboxes
+        )
+
+        footnote_bboxes_tmp = []
+        footnote_bboxes_tmp.extend(footnote_bboxes_by_model)
+        footnote_bboxes_tmp.extend(footnote_bboxes_by_rule)
+
+        page_info = construct_page_component(
+            page_id,
+            image_info,
+            table_info,
+            sorted_text_remain_text_block,
+            layout_bboxes,
+            inline_eq_info,
+            interline_eq_info,
+            page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"],
+            removed_text_blocks=removed_text_blocks,
+            removed_image_blocks=removed_images,
+            images_backup=images_backup,
+            droped_table_block=[],
+            table_backup=[],
+            layout_tree=layout_tree,
+            page_w=page.rect.width,
+            page_h=page.rect.height,
+            footnote_bboxes_tmp=footnote_bboxes_tmp,
+        )
+
+        page_info["image_bboxes_with_caption"] = image_bboxes_with_caption  # add by xr
+        page_info["table_bboxes_with_caption"] = table_bboxes_with_caption
+
+        page_info["bak_page_no_bboxes"] = page_no_bboxs
+        page_info["bak_header_bboxes"] = header_bboxs
+        page_info["bak_footer_bboxes"] = footer_bboxs
+
+        pdf_info_dict[f"page_{page_id}"] = page_info
+
+    # end page for
+
+    """计算后处理阶段耗时"""
+    start_time = time.time()
+
+    """
+    ==================================================================================================================================
+    去掉页眉和页脚，这里需要用到一定的统计量，所以放到最后
+    页眉和页脚主要从文本box和图片box中去除，位于页面的四周。
+    下面函数会直接修改pdf_info_dict,从文字块中、图片中删除属于页眉页脚的内容，删除内容做相对应记录
+    """
+    # 去页眉页脚
+    header, footer = drop_footer_header(
+        pdf_info_dict
+    )  # TODO: using header and footer boxes here !
+
+    """对单个layout内footnote和他下面的所有textbbox合并"""
+
+    for page_key, page_info in pdf_info_dict.items():
+        page_info = merge_footnote_blocks(page_info, main_text_font)
+        page_info = remove_footnote_blocks(page_info)
+        pdf_info_dict[page_key] = page_info
+
+    """进入pdf后置过滤器，去掉一些不合理的pdf"""
+
+    i = 0
+    for page_info in pdf_info_dict.values():
+        is_good_pdf, err = pdf_post_filter(page_info)
+        if not is_good_pdf:
+            logger.warning(f"page_id: {i}, drop this pdf: {book_name}, reason: {err}")
+            if not debug_mode:
+                return err
+        i += 1
+
+    if debug_mode:
+        params_file_save_path = join_path(
+            save_tmp_path, "md", book_name, "preproc_out.json"
+        )
+        page_draw_rect_save_path = join_path(
+            save_tmp_path, "md", book_name, "layout.pdf"
+        )
+        # dir_path = os.path.dirname(page_draw_rect_save_path)
+        # if not os.path.exists(dir_path):
+        #     # 如果目录不存在，创建它
+        #     os.makedirs(dir_path)
+
+        with open(params_file_save_path, "w", encoding="utf-8") as f:
+            json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
+        # 先检测本地 page_draw_rect_save_path 是否存在，如果存在则删除
+        if os.path.exists(page_draw_rect_save_path):
+            os.remove(page_draw_rect_save_path)
+        # 绘制bbox和layout到pdf
+        draw_bbox_on_page(pdf_docs, pdf_info_dict, page_draw_rect_save_path)
+        draw_layout_bbox_on_page(
+            pdf_docs, pdf_info_dict, header, footer, page_draw_rect_save_path
+        )
+
+    if debug_mode:
+        # 打印后处理阶段耗时
+        logger.info(f"post_processing_time: {get_delta_time(start_time)}")
+
+    """
+    ==================================================================================================================================
+    进入段落处理-2阶段
+    """
+
+    # 处理行内文字间距较大问题
+    pdf_info_dict = solve_inline_too_large_interval(pdf_info_dict)
+
+    start_time = time.time()
+
+    para_process_pipeline = ParaProcessPipeline()
+
+    def _deal_with_text_exception(error_info):
+        logger.warning(
+            f"page_id: {page_id}, drop this pdf: {book_name}, reason: {error_info}"
+        )
+        if error_info == denseSingleLineBlockException_msg:
+            logger.warning(
+                f"Drop this pdf: {book_name}, reason: {DropReason.DENSE_SINGLE_LINE_BLOCK}"
+            )
+            result = {
+                "need_drop": True,
+                "drop_reason": DropReason.DENSE_SINGLE_LINE_BLOCK,
+            }
+            return result
+        if error_info == titleDetectionException_msg:
+            logger.warning(
+                f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_DETECTION_FAILED}"
+            )
+            result = {
+                "need_drop": True,
+                "drop_reason": DropReason.TITLE_DETECTION_FAILED,
+            }
+            return result
+        elif error_info == titleLevelException_msg:
+            logger.warning(
+                f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_LEVEL_FAILED}"
+            )
+            result = {"need_drop": True, "drop_reason": DropReason.TITLE_LEVEL_FAILED}
+            return result
+        elif error_info == paraSplitException_msg:
+            logger.warning(
+                f"Drop this pdf: {book_name}, reason: {DropReason.PARA_SPLIT_FAILED}"
+            )
+            result = {"need_drop": True, "drop_reason": DropReason.PARA_SPLIT_FAILED}
+            return result
+        elif error_info == paraMergeException_msg:
+            logger.warning(
+                f"Drop this pdf: {book_name}, reason: {DropReason.PARA_MERGE_FAILED}"
+            )
+            result = {"need_drop": True, "drop_reason": DropReason.PARA_MERGE_FAILED}
+            return result
+
+    if debug_mode:
+        input_pdf_file = f"{pdf_local_path}.pdf"
+        output_dir = f"{save_path}/{book_name}"
+        output_pdf_file = f"{output_dir}/pdf_annos.pdf"
+
+        """
+        Call the para_process_pipeline function to process the pdf_info_dict.
+        
+        Parameters:
+        para_debug_mode: str or None
+            If para_debug_mode is None, the para_process_pipeline will not keep any intermediate results.
+            If para_debug_mode is "simple", the para_process_pipeline will only keep the annos on the pdf and the final results as a json file.
+            If para_debug_mode is "full", the para_process_pipeline will keep all the intermediate results generated during each step.
+        """
+        pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline(
+            pdf_info_dict,
+            para_debug_mode="simple",
+            input_pdf_path=input_pdf_file,
+            output_pdf_path=output_pdf_file,
+        )
+        # 打印段落处理阶段耗时
+        logger.info(f"para_process_time: {get_delta_time(start_time)}")
+
+        # debug的时候不return drop信息
+        if error_info is not None:
+            _deal_with_text_exception(error_info)
+        return pdf_info_dict
+    else:
+        pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline(
+            pdf_info_dict
+        )
+        if error_info is not None:
+            return _deal_with_text_exception(error_info)
+
+    return pdf_info_dict

magic_pdf/pipeline.py

@@ -3,61 +3,75 @@ 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.dict2md.ocr_mkcontent import (
+    ocr_mk_nlp_markdown,
+    ocr_mk_mm_markdown,
+    ocr_mk_mm_standard_format,
+    ocr_mk_mm_markdown_with_para,
+)
+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
+from magic_pdf.dict2md.mkcontent import mk_nlp_markdown, mk_universal_format
 from magic_pdf.pdf_parse_by_model import parse_pdf_by_model
 from magic_pdf.filter.pdf_classify_by_type import classify
 from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
 from loguru import logger
-
+from magic_pdf.pdf_parse_for_train import parse_pdf_for_train
+from magic_pdf.train_utils.convert_to_train_format import convert_to_train_format
 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):
     logger.exception(e)
-    jso['need_drop'] = True
-    jso['drop_reason'] = DropReason.Exception
-    jso['exception'] = f"ERROR: {e}"
+    jso["need_drop"] = True
+    jso["drop_reason"] = DropReason.Exception
+    jso["exception"] = f"ERROR: {e}"
     return jso
 
 
 def get_data_type(jso: dict):
-    data_type = jso.get('data_type')
+    data_type = jso.get("data_type")
     if data_type is None:
-        data_type = jso.get('file_type')
+        data_type = jso.get("file_type")
     return data_type
 
 
 def get_bookid(jso: dict):
-    book_id = jso.get('bookid')
+    book_id = jso.get("bookid")
     if book_id is None:
-        book_id = jso.get('original_file_id')
+        book_id = jso.get("original_file_id")
     return book_id
 
 
 def get_data_source(jso: dict):
-    data_source = jso.get('data_source')
+    data_source = jso.get("data_source")
     if data_source is None:
-        data_source = jso.get('file_source')
+        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_pdf_path = jso.get("file_location")
     s3_config = get_s3_config(s3_pdf_path)
     if doc_layout_check:
-        if 'doc_layout_result' not in jso:  # 检测json中是存在模型数据，如果没有则需要跳过该pdf
-            jso['need_drop'] = True
-            jso['drop_reason'] = DropReason.MISS_DOC_LAYOUT_RESULT
+        if (
+            "doc_layout_result" not in jso
+        ):  # 检测json中是存在模型数据，如果没有则需要跳过该pdf
+            jso["need_drop"] = True
+            jso["drop_reason"] = DropReason.MISS_DOC_LAYOUT_RESULT
             return jso
     try:
         data_source = get_data_source(jso)
-        file_id = jso.get('file_id')
-        book_name = data_source + "/" + file_id
+        file_id = jso.get("file_id")
+        book_name = f"{data_source}/{file_id}"
 
         # 首页存在超量drawing问题
         # special_pdf_list = ['zlib/zlib_21822650']
@@ -67,90 +81,111 @@ def meta_scan(jso: dict, doc_layout_check=True) -> dict:
         #     return jso
 
         start_time = time.time()  # 记录开始时间
-        logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
+        logger.info(
+            f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
+            file=sys.stderr,
+        )
         file_content = read_file(s3_pdf_path, s3_config)
         read_file_time = int(time.time() - start_time)  # 计算执行时间
 
         start_time = time.time()  # 记录开始时间
         res = pdf_meta_scan(s3_pdf_path, file_content)
-        if res.get('need_drop', False):  # 如果返回的字典里有need_drop，则提取drop_reason并跳过本次解析
-            jso['need_drop'] = True
-            jso['drop_reason'] = res["drop_reason"]
+        if res.get(
+            "need_drop", False
+        ):  # 如果返回的字典里有need_drop，则提取drop_reason并跳过本次解析
+            jso["need_drop"] = True
+            jso["drop_reason"] = res["drop_reason"]
         else:  # 正常返回
-            jso['pdf_meta'] = res
-            jso['content'] = ""
-            jso['remark'] = ""
-            jso['data_url'] = ""
+            jso["pdf_meta"] = res
+            jso["content"] = ""
+            jso["remark"] = ""
+            jso["data_url"] = ""
         end_time = time.time()  # 记录结束时间
         meta_scan_time = int(end_time - start_time)  # 计算执行时间
-        logger.info(f"book_name is:{book_name},end_time is:{formatted_time(end_time)},read_file_time is:{read_file_time},meta_scan_time is:{meta_scan_time}", file=sys.stderr)
-        jso['read_file_time'] = read_file_time
-        jso['meta_scan_time'] = meta_scan_time
+        logger.info(
+            f"book_name is:{book_name},end_time is:{formatted_time(end_time)},read_file_time is:{read_file_time},meta_scan_time is:{meta_scan_time}",
+            file=sys.stderr,
+        )
+        jso["read_file_time"] = read_file_time
+        jso["meta_scan_time"] = meta_scan_time
     except Exception as e:
         jso = exception_handler(jso, e)
     return jso
 
 
 def classify_by_type(jso: dict, debug_mode=False) -> dict:
-    #检测debug开关
+    # 检测debug开关
     if debug_mode:
         pass
-    else:# 如果debug没开，则检测是否有needdrop字段
-        if jso.get('need_drop', False):
+    else:  # 如果debug没开，则检测是否有needdrop字段
+        if jso.get("need_drop", False):
             return jso
     # 开始正式逻辑
     try:
-        pdf_meta = jso.get('pdf_meta')
+        pdf_meta = jso.get("pdf_meta")
         data_source = get_data_source(jso)
-        file_id = jso.get('file_id')
-        book_name = data_source + "/" + file_id
+        file_id = jso.get("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"]
         img_sz_list = pdf_meta["image_info_per_page"]
-        img_num_list = pdf_meta['imgs_per_page']
-        text_len_list = pdf_meta['text_len_per_page']
-        text_layout_list = pdf_meta['text_layout_per_page']
-        text_language = pdf_meta['text_language']
+        img_num_list = pdf_meta["imgs_per_page"]
+        text_len_list = pdf_meta["text_len_per_page"]
+        text_layout_list = pdf_meta["text_layout_per_page"]
+        text_language = pdf_meta["text_language"]
         # allow_language = ['zh', 'en']  # 允许的语言,目前只允许简中和英文的
 
         # if text_language not in allow_language:  # 如果语言不在允许的语言中，则drop
         #     jso['need_drop'] = True
         #     jso['drop_reason'] = DropReason.NOT_ALLOW_LANGUAGE
         #     return jso
-        pdf_path = pdf_meta['pdf_path']
-        is_encrypted = pdf_meta['is_encrypted']
-        is_needs_password = pdf_meta['is_needs_password']
-        if is_encrypted or is_needs_password:  # 加密的，需要密码的，没有页面的，都不处理
-            jso['need_drop'] = True
-            jso['drop_reason'] = DropReason.ENCRYPTED
+        pdf_path = pdf_meta["pdf_path"]
+        is_encrypted = pdf_meta["is_encrypted"]
+        is_needs_password = pdf_meta["is_needs_password"]
+        if (
+            is_encrypted or is_needs_password
+        ):  # 加密的，需要密码的，没有页面的，都不处理
+            jso["need_drop"] = True
+            jso["drop_reason"] = DropReason.ENCRYPTED
         else:
             start_time = time.time()  # 记录开始时间
-            is_text_pdf, results = classify(pdf_path, total_page, page_width, page_height, img_sz_list, text_len_list, img_num_list, text_layout_list)
+            is_text_pdf, results = classify(
+                pdf_path,
+                total_page,
+                page_width,
+                page_height,
+                img_sz_list,
+                text_len_list,
+                img_num_list,
+                text_layout_list,
+            )
             classify_time = int(time.time() - start_time)  # 计算执行时间
             if is_text_pdf:
-                pdf_meta['is_text_pdf'] = is_text_pdf
-                jso['pdf_meta'] = pdf_meta
-                jso['classify_time'] = classify_time
+                pdf_meta["is_text_pdf"] = is_text_pdf
+                jso["pdf_meta"] = pdf_meta
+                jso["classify_time"] = classify_time
                 # print(json.dumps(pdf_meta, ensure_ascii=False))
 
-                allow_language = ['zh', 'en']  # 允许的语言,目前只允许简中和英文的
-                if text_language not in allow_language:  # 如果语言不在允许的语言中，则drop
-                    jso['need_drop'] = True
-                    jso['drop_reason'] = DropReason.NOT_ALLOW_LANGUAGE
+                allow_language = ["zh", "en"]  # 允许的语言,目前只允许简中和英文的
+                if (
+                    text_language not in allow_language
+                ):  # 如果语言不在允许的语言中，则drop
+                    jso["need_drop"] = True
+                    jso["drop_reason"] = DropReason.NOT_ALLOW_LANGUAGE
                     return jso
             else:
                 # 先不drop
-                pdf_meta['is_text_pdf'] = is_text_pdf
-                jso['pdf_meta'] = pdf_meta
-                jso['classify_time'] = classify_time
-                jso['need_drop'] = True
-                jso['drop_reason'] = DropReason.NOT_IS_TEXT_PDF
+                pdf_meta["is_text_pdf"] = is_text_pdf
+                jso["pdf_meta"] = pdf_meta
+                jso["classify_time"] = classify_time
+                jso["need_drop"] = True
+                jso["drop_reason"] = DropReason.NOT_IS_TEXT_PDF
                 extra_info = {"classify_rules": []}
                 for condition, result in results.items():
                     if not result:
                         extra_info["classify_rules"].append(condition)
-                jso['extra_info'] = extra_info
+                jso["extra_info"] = extra_info
 
     except Exception as e:
         jso = exception_handler(jso, e)
@@ -161,48 +196,69 @@ def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:
 
     if debug_mode:
         pass
-    else:# 如果debug没开，则检测是否有needdrop字段
-        if jso.get('need_drop', False):
-            logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop", file=sys.stderr)
+    else:  # 如果debug没开，则检测是否有needdrop字段
+        if jso.get("need_drop", False):
+            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 = get_data_source(jso)
-        file_id = jso.get('file_id')
-        book_name = data_source + "/" + file_id
-        title = jso.get('title')
-        url_encode_title = quote(title, safe='')
-        if data_source != 'scihub':
+        file_id = jso.get("file_id")
+        book_name = f"{data_source}/{file_id}"
+        title = jso.get("title")
+        url_encode_title = quote(title, safe="")
+        if data_source != "scihub":
             return jso
-        pdf_intermediate_dict = jso['pdf_intermediate_dict']
+        pdf_intermediate_dict = jso["pdf_intermediate_dict"]
         # 将 pdf_intermediate_dict 解压
         pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
         i = 0
         for page in pdf_intermediate_dict.values():
-            if page.get('tables'):
-                if len(page['tables']) > 0:
+            if page.get("tables"):
+                if len(page["tables"]) > 0:
                     j = 0
-                    for table in page['tables']:
+                    for table in page["tables"]:
                         if debug_mode:
-                            image_path = join_path("s3://mllm-raw-media/pdf2md_img/", book_name, table['image_path'])
+                            image_path = join_path(
+                                "s3://mllm-raw-media/pdf2md_img/",
+                                book_name,
+                                table["image_path"],
+                            )
                         else:
-                            image_path = join_path("s3://mllm-raw-media/pdf2md_img/", table['image_path'])
+                            image_path = join_path(
+                                "s3://mllm-raw-media/pdf2md_img/", table["image_path"]
+                            )
 
-                        if image_path.endswith('.jpg'):
+                        if image_path.endswith(".jpg"):
                             j += 1
                             s3_client = get_s3_client(image_path)
                             bucket_name, bucket_key = parse_bucket_key(image_path)
                             # 通过s3_client获取图片到内存
-                            image_bytes = s3_client.get_object(Bucket=bucket_name, Key=bucket_key)['Body'].read()
+                            image_bytes = s3_client.get_object(
+                                Bucket=bucket_name, Key=bucket_key
+                            )["Body"].read()
                             # 保存图片到新的位置
                             if debug_mode:
-                                new_image_path = join_path("s3://mllm-raw-media/pdf2md_img/table_new/", url_encode_title + "_" + table['image_path'].lstrip('tables/'))
+                                new_image_path = join_path(
+                                    "s3://mllm-raw-media/pdf2md_img/table_new/",
+                                    url_encode_title
+                                    + "_"
+                                    + table["image_path"].lstrip("tables/"),
+                                )
                             else:
-                                new_image_path = join_path("s3://mllm-raw-media/pdf2md_img/table_new/", url_encode_title + f"_page{i}_{j}.jpg")
+                                new_image_path = join_path(
+                                    "s3://mllm-raw-media/pdf2md_img/table_new/",
+                                    url_encode_title + f"_page{i}_{j}.jpg",
+                                )
 
                             logger.info(new_image_path, file=sys.stderr)
                             bucket_name, bucket_key = parse_bucket_key(new_image_path)
-                            s3_client.put_object(Bucket=bucket_name, Key=bucket_key, Body=image_bytes)
+                            s3_client.put_object(
+                                Bucket=bucket_name, Key=bucket_key, Body=image_bytes
+                            )
                         else:
                             continue
             i += 1
@@ -217,8 +273,11 @@ 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:{get_data_source(jso)}/{jso['file_id']} need drop", file=sys.stderr)
+    if jso.get("need_drop", False):
+        logger.info(
+            f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop",
+            file=sys.stderr,
+        )
         jso["dropped"] = True
     return jso
 
@@ -227,19 +286,20 @@ def 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'])
+    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 = jso["pdf_intermediate_dict"]
         # 将 pdf_intermediate_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:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}", file=sys.stderr)
+        # markdown_content = mk_nlp_markdown(pdf_intermediate_dict)
+        jso["content_list"] = mk_universal_format(pdf_intermediate_dict)
+        # jso["content"] = markdown_content
+        logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']}")
         # 把无用的信息清空
         jso["doc_layout_result"] = ""
         jso["pdf_intermediate_dict"] = ""
@@ -250,19 +310,19 @@ def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
 
 
 def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
-    #检测debug开关
+    # 检测debug开关
     if debug_mode:
         pass
-    else:# 如果debug没开，则检测是否有needdrop字段
-        if jso.get('need_drop', False):
+    else:  # 如果debug没开，则检测是否有needdrop字段
+        if jso.get("need_drop", False):
             return jso
     # 开始正式逻辑
-    s3_pdf_path = jso.get('file_location')
+    s3_pdf_path = jso.get("file_location")
     s3_config = get_s3_config(s3_pdf_path)
-    model_output_json_list = jso.get('doc_layout_result')
+    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
+    file_id = jso.get("file_id")
+    book_name = f"{data_source}/{file_id}"
 
     # 1.23.22已修复
     # if debug_mode:
@@ -273,66 +333,114 @@ def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
     #         jso['drop_reason'] = DropReason.SPECIAL_PDF
     #         return jso
 
-    junk_img_bojids = jso['pdf_meta']['junk_img_bojids']
+    junk_img_bojids = jso["pdf_meta"]["junk_img_bojids"]
     # total_page = jso['pdf_meta']['total_page']
 
     # 增加检测 max_svgs 数量的检测逻辑，如果 max_svgs 超过3000则drop
-    svgs_per_page_list = jso['pdf_meta']['svgs_per_page']
+    svgs_per_page_list = jso["pdf_meta"]["svgs_per_page"]
     max_svgs = max(svgs_per_page_list)
     if max_svgs > 3000:
-        jso['need_drop'] = True
-        jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
+        jso["need_drop"] = True
+        jso["drop_reason"] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
     # elif total_page > 1000:
     #     jso['need_drop'] = True
     #     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和解析开始的时间
-            logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
-            pdf_info_dict = parse_pdf_by_model(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, junk_img_bojids=junk_img_bojids,
-                                                  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"]
+            logger.info(
+                f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
+                file=sys.stderr,
+            )
+            pdf_info_dict = parse_pdf_by_model(
+                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,
+                junk_img_bojids=junk_img_bojids,
+                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
+                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
+            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
 
 
+"""
+统一处理逻辑
+1.先调用parse_pdf对文本类pdf进行处理
+2.再调用ocr_dropped_parse_pdf,对之前drop的pdf进行处理
+"""
+
+
+def uni_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
+    jso = parse_pdf(jso, start_page_id=start_page_id, debug_mode=debug_mode)
+    jso = ocr_dropped_parse_pdf(jso, start_page_id=start_page_id, debug_mode=debug_mode)
+    return jso
+
+
+# 专门用来跑被drop的pdf，跑完之后需要把need_drop字段置为false
+def ocr_dropped_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
+    if not jso.get("need_drop", False):
+        return jso
+    else:
+        jso = ocr_parse_pdf_core(
+            jso, start_page_id=start_page_id, debug_mode=debug_mode
+        )
+        jso["need_drop"] = False
+        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):
+        if jso.get("need_drop", False):
             return jso
 
-    s3_pdf_path = jso.get('file_location')
+    jso = ocr_parse_pdf_core(jso, start_page_id=start_page_id, debug_mode=debug_mode)
+    return jso
+
+
+def ocr_parse_pdf_core(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')
+    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
+    file_id = jso.get("file_id")
+    book_name = f"{data_source}/{file_id}"
     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和解析开始的时间
-        logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
+        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,
@@ -342,20 +450,18 @@ def ocr_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
             pdf_model_profile=None,
             image_s3_config=image_s3_config,
             start_page_id=start_page_id,
-            debug_mode=debug_mode
+            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
+        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
+        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
@@ -366,18 +472,21 @@ 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'])
+        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 = 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)
+        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"] = ""
@@ -387,5 +496,145 @@ def ocr_pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
     return jso
 
 
+def ocr_pdf_intermediate_dict_to_markdown_with_para_for_qa(
+    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_with_para(pdf_intermediate_dict)
+        jso["content_ocr"] = 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"] = pdf_intermediate_dict
+        jso["pdf_meta"] = ""
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso
+
+
+def ocr_pdf_intermediate_dict_to_standard_format(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)
+        standard_format = ocr_mk_mm_standard_format(pdf_intermediate_dict)
+        jso["content_list"] = standard_format
+        logger.info(
+            f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",
+            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
+
+
+def parse_pdf_for_model_train(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}"
+
+    # 1.23.22已修复
+    # if debug_mode:
+    #     pass
+    # else:
+    #     if book_name == "zlib/zlib_21929367":
+    #         jso['need_drop'] = True
+    #         jso['drop_reason'] = DropReason.SPECIAL_PDF
+    #         return jso
+
+    junk_img_bojids = jso["pdf_meta"]["junk_img_bojids"]
+    # total_page = jso['pdf_meta']['total_page']
+
+    # 增加检测 max_svgs 数量的检测逻辑，如果 max_svgs 超过3000则drop
+    svgs_per_page_list = jso["pdf_meta"]["svgs_per_page"]
+    max_svgs = max(svgs_per_page_list)
+    if max_svgs > 3000:
+        jso["need_drop"] = True
+        jso["drop_reason"] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
+    # elif total_page > 1000:
+    #     jso['need_drop'] = True
+    #     jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
+    else:
+        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_for_train(
+                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,
+                junk_img_bojids=junk_img_bojids,
+                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 压缩并存储
+                jso["parsed_results"] = convert_to_train_format(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
+
+
 if __name__ == "__main__":
     pass

magic_pdf/pre_proc/equations_replace.py

@@ -6,9 +6,10 @@ import json
 import os
 from pathlib import Path
 from loguru import logger
+from magic_pdf.libs.ocr_content_type import ContentType
 
-TYPE_INLINE_EQUATION = "inline-equation"
-TYPE_INTERLINE_EQUATION = "interline-equation"
+TYPE_INLINE_EQUATION = ContentType.InlineEquation
+TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation
 
 
 def combine_chars_to_pymudict(block_dict, char_dict):

magic_pdf/pre_proc/ocr_cut_image.py

@@ -3,7 +3,7 @@ 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):
+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)
 
@@ -13,8 +13,8 @@ def cut_image_and_table(spans, page, page_id, book_name, save_path):
     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'))
+            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'))
+            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_dict_merge.py

@@ -2,6 +2,7 @@ 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.drop_tag import DropTag
 from magic_pdf.libs.ocr_content_type import ContentType
 
 
@@ -24,44 +25,48 @@ def line_sort_spans_by_left_to_right(lines):
     return line_objects
 
 def merge_spans_to_line(spans):
-    # 按照y0坐标排序
-    spans.sort(key=lambda span: span['bbox'][1])
+    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 = []
+        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)
-            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
+        return lines
 
 def merge_spans_to_line_by_layout(spans, layout_bboxes):
     lines = []
     new_spans = []
+    dropped_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:
+            if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], layout_bbox) > 0.6:
                 layout_sapns.append(span)
         # 如果layout_sapns不为空，则放入new_spans中
         if len(layout_sapns) > 0:
@@ -75,10 +80,14 @@ def merge_spans_to_line_by_layout(spans, layout_bboxes):
             layout_lines = merge_spans_to_line(layout_sapns)
             lines.extend(layout_lines)
 
-    #对line中的span进行排序
+    # 对line中的span进行排序
     lines = line_sort_spans_by_left_to_right(lines)
 
-    return lines
+    for span in spans:
+        span['tag'] = DropTag.NOT_IN_LAYOUT
+        dropped_spans.append(span)
+
+    return lines, dropped_spans
 
 
 def merge_lines_to_block(lines):

magic_pdf/pre_proc/ocr_span_list_modify.py

@@ -2,10 +2,12 @@ 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.drop_tag import DropTag
 from magic_pdf.libs.ocr_content_type import ContentType
 
 
 def remove_overlaps_min_spans(spans):
+    dropped_spans = []
     #  删除重叠spans中较小的那些
     for span1 in spans.copy():
         for span2 in spans.copy():
@@ -15,7 +17,9 @@ def remove_overlaps_min_spans(spans):
                     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
+                        bbox_to_remove['tag'] = DropTag.SPAN_OVERLAP
+                        dropped_spans.append(bbox_to_remove)
+    return spans, dropped_spans
 
 
 def remove_spans_by_bboxes(spans, need_remove_spans_bboxes):
@@ -35,9 +39,7 @@ def remove_spans_by_bboxes(spans, need_remove_spans_bboxes):
 
 
 def remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict):
-    dropped_text_block = []
-    dropped_image_block = []
-    dropped_table_block = []
+    dropped_spans = []
     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 = []
@@ -50,14 +52,9 @@ def remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict):
         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)
+            dropped_spans.append(span)
 
-    return spans, dropped_text_block, dropped_image_block, dropped_table_block
+    return spans, dropped_spans
 
 
 def adjust_bbox_for_standalone_block(spans):
@@ -77,70 +74,74 @@ def adjust_bbox_for_standalone_block(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])
 
-    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])
 
-    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
 
-    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)
-            # 则开始新行
+            # 如果当前的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)))
-            current_line = [span]
-            line_first_y0 = span["bbox"][1]
-            line_first_y = span["bbox"][3]
-            continue
+        for line in text_inline_lines:
+            # 按照x0坐标排序
+            current_line = line[0]
+            current_line.sort(key=lambda span: span['bbox'][0])
 
-        # 如果当前的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)
+        # 调整每一个文字行内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
 
-        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
+        # return spans, displayed_list, text_inline_lines
 
 
 def modify_inline_equation(spans: list, displayed_list: list, text_inline_lines: list):
@@ -157,7 +158,7 @@ def modify_inline_equation(spans: list, displayed_list: list, text_inline_lines:
             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)):
+                span['bbox'], (0, y0, 0, y1)):
 
                 # 调整公式类型
                 if span["type"] == ContentType.InterlineEquation:

magic_pdf/pre_proc/remove_footer_header.py

@@ -1,6 +1,7 @@
 import re
 
 from magic_pdf.libs.boxbase import _is_in_or_part_overlap
+from magic_pdf.libs.drop_tag import CONTENT_IN_FOOT_OR_HEADER, PAGE_NO
 
 
 def remove_headder_footer_one_page(text_raw_blocks, image_bboxes, table_bboxes, header_bboxs, footer_bboxs,
@@ -67,7 +68,7 @@ def remove_headder_footer_one_page(text_raw_blocks, image_bboxes, table_bboxes,
                 blk['lines'].remove(line)
         else:
             # if not blk['lines']:
-            blk['tag'] = 'in-foot-header-area'
+            blk['tag'] = CONTENT_IN_FOOT_OR_HEADER
             text_block_to_remove.append(blk)
 
     """有的时候由于pageNo太小了，总是会有一点和content_boundry重叠一点，被放入正文，因此对于pageNo，进行span粒度的删除"""
@@ -80,7 +81,7 @@ def remove_headder_footer_one_page(text_raw_blocks, image_bboxes, table_bboxes,
                         for span in line['spans']:
                             if _is_in_or_part_overlap(pagenobox, span['bbox']):
                                 # span['text'] = ''
-                                span['tag'] = "page-no"
+                                span['tag'] = PAGE_NO
                                 # 检查这个block是否只有这一个span，如果是，那么就把这个block也删除
                                 if len(line['spans']) == 1 and len(block['lines']) == 1:
                                     page_no_block_2_remove.append(block)
@@ -96,7 +97,7 @@ def remove_headder_footer_one_page(text_raw_blocks, image_bboxes, table_bboxes,
                     if last_span['text'].strip() and not re.search('[a-zA-Z]', last_span['text']) and re.search('[0-9]',
                                                                                                                 last_span[
                                                                                                                     'text']):
-                        last_span['tag'] = "page-no"
+                        last_span['tag'] = PAGE_NO
                         page_no_block_2_remove.append(last_block)
 
     for b in page_no_block_2_remove:

magic_pdf/pre_proc/remove_rotate_bbox.py

@@ -1,6 +1,7 @@
 import math
 
 from magic_pdf.libs.boxbase import is_vbox_on_side
+from magic_pdf.libs.drop_tag import EMPTY_SIDE_BLOCK, ROTATE_TEXT, VERTICAL_TEXT
 
 
 def detect_non_horizontal_texts(result_dict):
@@ -134,13 +135,13 @@ def remove_rotate_side_textblock(pymu_text_block, page_width, page_height):
             is_box_valign = (len(set([int(line['spans'][0]['bbox'][0] ) for line in lines if len(line['spans'])>0]))==1) and (len([int(line['spans'][0]['bbox'][0] ) for line in lines if len(line['spans'])>0])>1)  # 测试bbox在垂直方向是不是x0都相等，也就是在垂直方向排列.同时必须大于等于2个字
             
             if is_box_valign:
-                block['tag'] = "vertical-text"
+                block['tag'] = VERTICAL_TEXT
                 removed_text_block.append(block)
                 continue
         
         for line in lines:
             if line['dir']!=(1,0):
-                block['tag'] = "rotate"
+                block['tag'] = ROTATE_TEXT
                 removed_text_block.append(block) # 只要有一个line不是dir=(1,0)，就把整个block都删掉
                 break
         
@@ -177,7 +178,7 @@ def remove_side_blank_block(pymu_text_block, page_width, page_height):
            continue
             
         if __is_empty_side_box(block):
-            block['tag'] = "empty-side-block"
+            block['tag'] = EMPTY_SIDE_BLOCK
             removed_text_block.append(block)
             continue
         

magic_pdf/pre_proc/resolve_bbox_conflict.py

@@ -6,6 +6,7 @@
 """
 
 from magic_pdf.libs.boxbase import _is_in, _is_in_or_part_overlap, _is_left_overlap
+from magic_pdf.libs.drop_tag import ON_IMAGE_TEXT, ON_TABLE_TEXT
 
 
 def resolve_bbox_overlap_conflict(images:list, tables:list, interline_equations:list, inline_equations:list, text_raw_blocks:list):
@@ -27,14 +28,14 @@ def resolve_bbox_overlap_conflict(images:list, tables:list, interline_equations:
         for text_block in text_raw_blocks:
             text_bbox = text_block["bbox"]
             if _is_in(text_bbox, image_box):
-                text_block['tag'] = "on-image"
+                text_block['tag'] = ON_IMAGE_TEXT
                 text_block_removed.append(text_block)
     # 去掉table上的文字block
     for table_box in tables:
         for text_block in text_raw_blocks:
             text_bbox = text_block["bbox"]
             if _is_in(text_bbox, table_box):
-                text_block['tag'] = "on-table"
+                text_block['tag'] = ON_TABLE_TEXT
                 text_block_removed.append(text_block)
                 
     for text_block in text_block_removed:

magic_pdf/train_utils/convert_to_train_format.py

@@ -0,0 +1,53 @@
+def convert_to_train_format(jso: dict) -> []:
+    pages = []
+    for k, v in jso.items():
+        if not k.startswith("page_"):
+            continue
+        page_idx = v["page_idx"]
+        width, height = v["page_size"]
+
+        info = {"page_info": {"page_no": page_idx, "height": height, "width": width}}
+
+        bboxes: list[dict] = []
+        for img_bbox in v["image_bboxes_with_caption"]:
+            bbox = {"category_id": 1, "bbox": img_bbox["bbox"]}
+            if "caption" in img_bbox:
+                bbox["caption_bbox"] = img_bbox["caption"]
+            bboxes.append(bbox)
+
+        for tbl_bbox in v["table_bboxes_with_caption"]:
+            bbox = {"category_id": 7, "bbox": tbl_bbox["bbox"]}
+            if "caption" in tbl_bbox:
+                bbox["caption_bbox"] = tbl_bbox["caption"]
+            bboxes.append(bbox)
+
+        for bbox in v["bak_page_no_bboxes"]:
+            n_bbox = {"category_id": 4, "bbox": bbox}
+            bboxes.append(n_bbox)
+
+        for bbox in v["bak_header_bboxes"]:
+            n_bbox = {"category_id": 3, "bbox": bbox}
+            bboxes.append(n_bbox)
+
+        for bbox in v["bak_footer_bboxes"]:
+            n_bbox = {"category_id": 6, "bbox": bbox}
+            bboxes.append(n_bbox)
+
+        # 脚注， 目前没有看到例子
+        for para in v["para_blocks"]:
+            n_bbox = {"category_id": 2, "bbox": para["bbox"]}
+            bboxes.append(n_bbox)
+
+        for inline_equation in v["inline_equations"]:
+            n_bbox = {"category_id": 13, "bbox": inline_equation["bbox"]}
+            bboxes.append(n_bbox)
+
+        for inter_equation in v["interline_equations"]:
+            n_bbox = {"category_id": 10, "bbox": inter_equation["bbox"]}
+            bboxes.append(n_bbox)
+
+        info["bboxes"] = bboxes
+        info["layout_tree"] = v["layout_bboxes"]
+        pages.append(info)
+
+    return pages

magic_pdf/train_utils/extract_caption.py

@@ -0,0 +1,59 @@
+from magic_pdf.libs.boxbase import _is_in
+
+
+def extract_caption_bbox(outer: list, inner: list) -> list:
+    """
+    ret: list of {
+                    "bbox": [1,2,3,4],
+                    "caption": [5,6,7,8] # may existed
+                }
+
+    """
+    found_count = 0  # for debug
+    print(outer, inner)
+
+    def is_float_equal(a, b):
+        if 0.01 > abs(a - b):  # non strict float equal compare
+            return True
+        return False
+
+    outer_h = {i: outer[i] for i in range(len(outer))}
+    ret = []
+    for v in inner:
+        ix0, iy0, ix1, iy1 = v
+        found_idx = None
+        d = {"bbox": v[:4]}
+        for k in outer_h:
+            ox0, oy0, ox1, oy1 = outer_h[k]
+            equal_float_flags = [
+                is_float_equal(ix0, ox0),
+                is_float_equal(iy0, oy0),
+                is_float_equal(ix1, ox1),
+                is_float_equal(iy1, oy1),
+            ]
+            if _is_in(v, outer_h[k]) and not all(equal_float_flags):
+                found_idx = k
+                break
+        if found_idx is not None:
+            found_count += 1
+            captions: list[list] = []
+            ox0, oy0, ox1, oy1 = outer_h[found_idx]
+            captions = [
+                [ox0, oy0, ix0, oy1],
+                [ox0, oy0, ox1, iy0],
+                [ox0, iy1, ox1, oy1],
+                [ix1, oy0, ox1, oy1],
+            ]
+            captions = sorted(
+                captions,
+                key=lambda rect: abs(rect[0] - rect[2]) * abs(rect[1] - rect[3]),
+            )  # 面积最大的框就是caption
+            d["caption"] = captions[-1]
+            outer_h.pop(
+                found_idx
+            )  # 同一个 outer box 只能用于确定一个 inner box 的 caption 位置。
+
+        ret.append(d)
+
+    print("found_count: ", found_count)
+    return ret

magic_pdf/train_utils/remove_footer_header.py

@@ -0,0 +1,159 @@
+import re
+
+from magic_pdf.libs.boxbase import _is_in_or_part_overlap
+from magic_pdf.libs.drop_tag import CONTENT_IN_FOOT_OR_HEADER, PAGE_NO
+
+
+"""
+    copy from pre_proc/remove_footer_header.py
+"""
+
+
+def remove_headder_footer_one_page(
+    text_raw_blocks,
+    image_bboxes,
+    table_bboxes,
+    header_bboxs,
+    footer_bboxs,
+    page_no_bboxs,
+    page_w,
+    page_h,
+):
+    """
+    删除页眉页脚，页码
+    从line级别进行删除，删除之后观察这个text-block是否是空的，如果是空的，则移动到remove_list中
+    """
+    if 1:
+        return image_bboxes, table_bboxes, text_raw_blocks, [], [], []
+
+    header = []
+    footer = []
+    if len(header) == 0:
+        model_header = header_bboxs
+        if model_header:
+            x0 = min([x for x, _, _, _ in model_header])
+            y0 = min([y for _, y, _, _ in model_header])
+            x1 = max([x1 for _, _, x1, _ in model_header])
+            y1 = max([y1 for _, _, _, y1 in model_header])
+            header = [x0, y0, x1, y1]
+    if len(footer) == 0:
+        model_footer = footer_bboxs
+        if model_footer:
+            x0 = min([x for x, _, _, _ in model_footer])
+            y0 = min([y for _, y, _, _ in model_footer])
+            x1 = max([x1 for _, _, x1, _ in model_footer])
+            y1 = max([y1 for _, _, _, y1 in model_footer])
+            footer = [x0, y0, x1, y1]
+
+    header_y0 = 0 if len(header) == 0 else header[3]
+    footer_y0 = page_h if len(footer) == 0 else footer[1]
+    if page_no_bboxs:
+        top_part = [b for b in page_no_bboxs if b[3] < page_h / 2]
+        btn_part = [b for b in page_no_bboxs if b[1] > page_h / 2]
+
+        top_max_y0 = max([b[1] for b in top_part]) if top_part else 0
+        btn_min_y1 = min([b[3] for b in btn_part]) if btn_part else page_h
+
+        header_y0 = max(header_y0, top_max_y0)
+        footer_y0 = min(footer_y0, btn_min_y1)
+
+    content_boundry = [0, header_y0, page_w, footer_y0]
+
+    header = [0, 0, page_w, header_y0]
+    footer = [0, footer_y0, page_w, page_h]
+
+    """以上计算出来了页眉页脚的边界，下面开始进行删除"""
+    text_block_to_remove = []
+    # 首先检查每个textblock
+    for blk in text_raw_blocks:
+        if len(blk["lines"]) > 0:
+            for line in blk["lines"]:
+                line_del = []
+                for span in line["spans"]:
+                    span_del = []
+                    if span["bbox"][3] < header_y0:
+                        span_del.append(span)
+                    elif _is_in_or_part_overlap(
+                        span["bbox"], header
+                    ) or _is_in_or_part_overlap(span["bbox"], footer):
+                        span_del.append(span)
+                for span in span_del:
+                    line["spans"].remove(span)
+                if not line["spans"]:
+                    line_del.append(line)
+
+            for line in line_del:
+                blk["lines"].remove(line)
+        else:
+            # if not blk['lines']:
+            blk["tag"] = CONTENT_IN_FOOT_OR_HEADER
+            text_block_to_remove.append(blk)
+
+    """有的时候由于pageNo太小了，总是会有一点和content_boundry重叠一点，被放入正文，因此对于pageNo，进行span粒度的删除"""
+    page_no_block_2_remove = []
+    if page_no_bboxs:
+        for pagenobox in page_no_bboxs:
+            for block in text_raw_blocks:
+                if _is_in_or_part_overlap(
+                    pagenobox, block["bbox"]
+                ):  # 在span级别删除页码
+                    for line in block["lines"]:
+                        for span in line["spans"]:
+                            if _is_in_or_part_overlap(pagenobox, span["bbox"]):
+                                # span['text'] = ''
+                                span["tag"] = PAGE_NO
+                                # 检查这个block是否只有这一个span，如果是，那么就把这个block也删除
+                                if len(line["spans"]) == 1 and len(block["lines"]) == 1:
+                                    page_no_block_2_remove.append(block)
+    else:
+        # 测试最后一个是不是页码：规则是，最后一个block仅有1个line,一个span,且text是数字，空格，符号组成，不含字母,并且包含数字
+        if len(text_raw_blocks) > 0:
+            text_raw_blocks.sort(key=lambda x: x["bbox"][1], reverse=True)
+            last_block = text_raw_blocks[0]
+            if len(last_block["lines"]) == 1:
+                last_line = last_block["lines"][0]
+                if len(last_line["spans"]) == 1:
+                    last_span = last_line["spans"][0]
+                    if (
+                        last_span["text"].strip()
+                        and not re.search("[a-zA-Z]", last_span["text"])
+                        and re.search("[0-9]", last_span["text"])
+                    ):
+                        last_span["tag"] = PAGE_NO
+                        page_no_block_2_remove.append(last_block)
+
+    for b in page_no_block_2_remove:
+        text_block_to_remove.append(b)
+
+    for blk in text_block_to_remove:
+        if blk in text_raw_blocks:
+            text_raw_blocks.remove(blk)
+
+    text_block_remain = text_raw_blocks
+    image_bbox_to_remove = [
+        bbox
+        for bbox in image_bboxes
+        if not _is_in_or_part_overlap(bbox, content_boundry)
+    ]
+
+    image_bbox_remain = [
+        bbox for bbox in image_bboxes if _is_in_or_part_overlap(bbox, content_boundry)
+    ]
+    table_bbox_to_remove = [
+        bbox
+        for bbox in table_bboxes
+        if not _is_in_or_part_overlap(bbox, content_boundry)
+    ]
+    table_bbox_remain = [
+        bbox for bbox in table_bboxes if _is_in_or_part_overlap(bbox, content_boundry)
+    ]
+
+    #        1,                 2,                3
+    return (
+        image_bbox_remain,
+        table_bbox_remain,
+        text_block_remain,
+        text_block_to_remove,
+        image_bbox_to_remove,
+        table_bbox_to_remove,
+    )

magic_pdf/train_utils/vis_utils.py

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

requirements.txt

@@ -11,5 +11,7 @@ pycld2>=0.41
 regex>=2023.12.25
 spacy>=3.7.4
 termcolor>=2.4.0
+scikit-learn
+wordninja>=2.0.0
 en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl
 zh_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/zh_core_web_sm-3.7.0/zh_core_web_sm-3.7.0-py3-none-any.whl

tests/overall_indicator.py

@@ -0,0 +1,462 @@
+import json
+import pandas as pd
+import numpy as np
+import re
+from nltk.translate.bleu_score import sentence_bleu
+import time
+import argparse
+import os
+from sklearn.metrics import classification_report,confusion_matrix
+from collections import Counter
+from sklearn import metrics
+from pandas import isnull
+
+
+def indicator_cal(json_standard,json_test):
+
+    json_standard = pd.DataFrame(json_standard)
+    json_test = pd.DataFrame(json_test)
+
+
+
+    '''数据集总体指标'''
+    
+    a=json_test[['id','mid_json']]
+    b=json_standard[['id','mid_json','pass_label']]
+    outer_merge=pd.merge(a,b,on='id',how='outer')
+    outer_merge.columns=['id','standard_mid_json','test_mid_json','pass_label']
+    standard_exist=outer_merge.standard_mid_json.apply(lambda x: not isnull(x))
+    test_exist=outer_merge.test_mid_json.apply(lambda x: not isnull(x))
+
+    overall_report = {}
+    overall_report['accuracy']=metrics.accuracy_score(standard_exist,test_exist)
+    overall_report['precision']=metrics.precision_score(standard_exist,test_exist)
+    overall_report['recall']=metrics.recall_score(standard_exist,test_exist)
+    overall_report['f1_score']=metrics.f1_score(standard_exist,test_exist)
+
+
+    inner_merge=pd.merge(a,b,on='id',how='inner')
+    inner_merge.columns=['id','standard_mid_json','test_mid_json','pass_label']
+    json_standard = inner_merge['standard_mid_json']#check一下是否对齐
+    json_test = inner_merge['test_mid_json']
+
+
+    
+
+    '''批量读取中间生成的json文件'''
+    test_inline_equations=[]
+    test_interline_equations=[]
+    test_dropped_text_bboxes=[]
+    test_dropped_text_tag=[]
+    test_dropped_image_bboxes=[]
+    test_dropped_table_bboxes=[] 
+    test_preproc_num=[]#阅读顺序
+    test_para_num=[]
+    test_para_text=[]
+
+    for i in json_test:
+        mid_json=pd.DataFrame(i)
+        mid_json=mid_json.iloc[:,:-1]
+        for j1 in mid_json.loc['inline_equations',:]:
+            page_in=[]
+            for k1 in j1:
+                page_in.append(k1['latex_text'])
+            test_inline_equations.append(page_in)
+        for j2 in mid_json.loc['interline_equations',:]:
+            page_in=[]
+            for k2 in j2:
+                page_in.append(k2['latex_text'])
+            test_interline_equations.append(page_in)
+
+        for j3 in mid_json.loc['droped_text_block',:]:
+            page_in_bbox=[]
+            page_in_tag=[]
+            for k3 in j3:
+                page_in_bbox.append(k3['bbox'])
+                #如果k3中存在tag这个key
+                if 'tag' in k3.keys():
+                    page_in_tag.append(k3['tag'])
+                else:
+                    page_in_tag.append('None')
+            test_dropped_text_tag.append(page_in_tag)
+            test_dropped_text_bboxes.append(page_in_bbox)
+        for j4 in mid_json.loc['droped_image_block',:]:
+                test_dropped_image_bboxes.append(j4)
+        for j5 in mid_json.loc['droped_table_block',:]:
+                test_dropped_table_bboxes.append(j5)
+        for j6 in mid_json.loc['preproc_blocks',:]:
+            page_in=[]
+            for k6 in j6:
+                page_in.append(k6['number'])
+            test_preproc_num.append(page_in)
+
+        test_pdf_text=[]     
+        for j7 in mid_json.loc['para_blocks',:]:
+            test_para_num.append(len(j7))  
+            for k7 in j7:
+                test_pdf_text.append(k7['text'])  
+        test_para_text.append(test_pdf_text)
+
+
+
+    standard_inline_equations=[]
+    standard_interline_equations=[]
+    standard_dropped_text_bboxes=[]
+    standard_dropped_text_tag=[]
+    standard_dropped_image_bboxes=[]
+    standard_dropped_table_bboxes=[] 
+    standard_preproc_num=[]#阅读顺序
+    standard_para_num=[]
+    standard_para_text=[]
+
+    for i in json_standard:
+        mid_json=pd.DataFrame(i)
+        mid_json=mid_json.iloc[:,:-1]
+        for j1 in mid_json.loc['inline_equations',:]:
+            page_in=[]
+            for k1 in j1:
+                page_in.append(k1['latex_text'])
+            standard_inline_equations.append(page_in)
+        for j2 in mid_json.loc['interline_equations',:]:
+            page_in=[]
+            for k2 in j2:
+                page_in.append(k2['latex_text'])
+            standard_interline_equations.append(page_in)
+        for j3 in mid_json.loc['droped_text_block',:]:
+            page_in_bbox=[]
+            page_in_tag=[]
+            for k3 in j3:
+                page_in_bbox.append(k3['bbox'])
+                if 'tag' in k3.keys():
+                    page_in_tag.append(k3['tag'])
+                else:
+                    page_in_tag.append('None')
+            standard_dropped_text_bboxes.append(page_in_bbox)
+            standard_dropped_text_tag.append(page_in_tag)
+        for j4 in mid_json.loc['droped_image_block',:]:
+                standard_dropped_image_bboxes.append(j4)
+        for j5 in mid_json.loc['droped_table_block',:]:
+                standard_dropped_table_bboxes.append(j5)
+        for j6 in mid_json.loc['preproc_blocks',:]:
+            page_in=[]
+            for k6 in j6:
+                page_in.append(k6['number'])
+            standard_preproc_num.append(page_in)     
+
+        standard_pdf_text=[]
+        for j7 in mid_json.loc['para_blocks',:]:
+            standard_para_num.append(len(j7))  
+            for k7 in j7:
+                standard_pdf_text.append(k7['text'])
+        standard_para_text.append(standard_pdf_text)
+
+
+    """
+    在计算指标之前最好先确认基本统计信息是否一致
+    """
+
+
+    '''
+    计算pdf之间的总体编辑距离和bleu
+    这里只计算正例的pdf
+    '''
+    
+    test_para_text=np.asarray(test_para_text, dtype = object)[inner_merge['pass_label']=='yes']
+    standard_para_text=np.asarray(standard_para_text, dtype = object)[inner_merge['pass_label']=='yes']
+
+    pdf_dis=[]
+    pdf_bleu=[]
+    for a,b in zip(test_para_text,standard_para_text):
+        a1=[ ''.join(i) for i in a]
+        b1=[ ''.join(i) for i in b]
+        pdf_dis.append(Levenshtein_Distance(a1,b1))
+        pdf_bleu.append(sentence_bleu([a1],b1))
+    overall_report['pdf间的平均编辑距离']=np.mean(pdf_dis)
+    overall_report['pdf间的平均bleu']=np.mean(pdf_bleu)
+
+
+    '''行内公式编辑距离和bleu'''
+    dis1=[]
+    bleu1=[]
+
+    test_inline_equations=[ ''.join(i) for i in test_inline_equations]
+    standard_inline_equations=[ ''.join(i) for i in standard_inline_equations]
+           
+    for a,b in zip(test_inline_equations,standard_inline_equations):
+        if len(a)==0 and len(b)==0:
+            continue
+        else:
+            if a==b:
+                dis1.append(0)
+                bleu1.append(1)
+            else:
+                dis1.append(Levenshtein_Distance(a,b))
+                bleu1.append(sentence_bleu([a],b))
+    inline_equations_edit=np.mean(dis1)
+    inline_equations_bleu=np.mean(bleu1)
+
+
+    '''行间公式编辑距离和bleu'''
+    dis2=[]
+    bleu2=[]
+
+    test_interline_equations=[ ''.join(i) for i in test_interline_equations]
+    standard_interline_equations=[ ''.join(i) for i in standard_interline_equations]
+
+    for a,b in zip(test_interline_equations,standard_interline_equations):
+        if len(a)==0 and len(b)==0:
+            continue
+        else:
+            if a==b:
+                dis2.append(0)
+                bleu2.append(1)
+            else:
+                dis2.append(Levenshtein_Distance(a,b))
+                bleu2.append(sentence_bleu([a],b))
+    interline_equations_edit=np.mean(dis2)
+    interline_equations_bleu=np.mean(bleu2)
+
+
+
+
+    '''可以先检查page和bbox数量是否一致'''
+
+    '''dropped_text_block的bbox匹配相关指标'''
+    test_text_bbox=[]
+    standard_text_bbox=[]
+    test_tag=[]
+    standard_tag=[]
+
+
+    index=0
+    for a,b in zip(test_dropped_text_bboxes,standard_dropped_text_bboxes):
+        test_page_tag=[]
+        standard_page_tag=[]
+        test_page_bbox=[]
+        standard_page_bbox=[]
+        if len(a)==0 and len(b)==0:
+            pass
+        else:
+            for i in range(len(b)):
+                judge=0
+                standard_page_tag.append(standard_dropped_text_tag[index][i])
+                standard_page_bbox.append(1)
+                for j in range(len(a)):
+                    if bbox_offset(b[i],a[j]):
+                        judge=1
+                        test_page_tag.append(test_dropped_text_tag[index][j])
+                        test_page_bbox.append(1)
+                        break
+                if judge==0:
+                    test_page_tag.append('None')
+                    test_page_bbox.append(0)
+
+
+            if len(test_dropped_text_tag[index])+test_page_tag.count('None')>len(standard_dropped_text_tag[index]):#有多删的情况出现
+                test_page_tag1=test_page_tag.copy()
+                if 'None' in test_page_tag:
+                    test_page_tag1=test_page_tag1.remove('None')
+                else:
+                    test_page_tag1=test_page_tag
+
+                diff=list((Counter(test_dropped_text_tag[index]) - Counter(test_page_tag1)).elements())
+              
+                test_page_tag.extend(diff)
+                standard_page_tag.extend(['None']*len(diff))
+                test_page_bbox.extend([1]*len(diff))
+                standard_page_bbox.extend([0]*len(diff))
+
+            test_tag.extend(test_page_tag)
+            standard_tag.extend(standard_page_tag)
+            test_text_bbox.extend(test_page_bbox)
+            standard_text_bbox.extend(standard_page_bbox)
+
+        index+=1
+
+    
+    text_block_report = {}
+    text_block_report['accuracy']=metrics.accuracy_score(standard_text_bbox,test_text_bbox)
+    text_block_report['precision']=metrics.precision_score(standard_text_bbox,test_text_bbox)
+    text_block_report['recall']=metrics.recall_score(standard_text_bbox,test_text_bbox)
+    text_block_report['f1_score']=metrics.f1_score(standard_text_bbox,test_text_bbox)
+
+    '''删除的text_block的tag的准确率,召回率和f1-score'''
+    text_block_tag_report = classification_report(y_true=standard_tag , y_pred=test_tag,output_dict=True)
+    del text_block_tag_report['None']
+    del text_block_tag_report["macro avg"]
+    del text_block_tag_report["weighted avg"]
+
+
+    '''dropped_image_block的bbox匹配相关指标'''
+    '''有数据格式不一致的问题'''
+    
+    test_image_bbox=[]
+    standard_image_bbox=[]
+    for a,b in zip(test_dropped_image_bboxes,standard_dropped_image_bboxes):
+
+        test_page_bbox=[]
+        standard_page_bbox=[]
+        if len(a)==0 and len(b)==0:
+            pass
+        else:
+            for i in b:
+                if len(i)!=4:
+                    continue
+                else:
+                    judge=0
+                    standard_page_bbox.append(1)
+                    for j in a:
+                        if bbox_offset(i,j):
+                            judge=1
+                            test_page_bbox.append(1)
+                            break
+                    if judge==0:
+                        test_page_bbox.append(0)
+                        
+            diff_num=len(a)+test_page_bbox.count(0)-len(b)
+            if diff_num>0:#有多删的情况出现
+                test_page_bbox.extend([1]*diff_num)
+                standard_page_bbox.extend([0]*diff_num)
+
+          
+            test_image_bbox.extend(test_page_bbox)
+            standard_image_bbox.extend(standard_page_bbox)
+
+    
+    image_block_report = {}
+    image_block_report['accuracy']=metrics.accuracy_score(standard_image_bbox,test_image_bbox)
+    image_block_report['precision']=metrics.precision_score(standard_image_bbox,test_image_bbox)
+    image_block_report['recall']=metrics.recall_score(standard_image_bbox,test_image_bbox)
+    image_block_report['f1_score']=metrics.f1_score(standard_image_bbox,test_image_bbox)
+    
+
+
+    
+    '''dropped_table_block的bbox匹配相关指标'''
+    test_table_bbox=[]
+    standard_table_bbox=[]
+    for a,b in zip(test_dropped_table_bboxes,standard_dropped_table_bboxes):
+
+        test_page_bbox=[]
+        standard_page_bbox=[]
+        if len(a)==0 and len(b)==0:
+            pass
+        else:
+            for i in b:
+                if len(i)!=4:
+                    continue
+                else:
+                    judge=0
+                    standard_page_bbox.append(1)
+                    for j in a:
+                        if bbox_offset(i,j):
+                            judge=1
+                            test_page_bbox.append(1)
+                            break
+                    if judge==0:
+                        test_page_bbox.append(0)
+                        
+            diff_num=len(a)+test_page_bbox.count(0)-len(b)
+            if diff_num>0:#有多删的情况出现
+                test_page_bbox.extend([1]*diff_num)
+                standard_page_bbox.extend([0]*diff_num)
+
+          
+            test_table_bbox.extend(test_page_bbox)
+            standard_table_bbox.extend(standard_page_bbox)
+
+    table_block_report = {}
+    table_block_report['accuracy']=metrics.accuracy_score(standard_table_bbox,test_table_bbox)
+    table_block_report['precision']=metrics.precision_score(standard_table_bbox,test_table_bbox)
+    table_block_report['recall']=metrics.recall_score(standard_table_bbox,test_table_bbox)
+    table_block_report['f1_score']=metrics.f1_score(standard_table_bbox,test_table_bbox)
+    
+   
+    '''阅读顺序编辑距离的均值'''
+    preproc_num_dis=[]
+    for a,b in zip(test_preproc_num,standard_preproc_num):
+        preproc_num_dis.append(Levenshtein_Distance(a,b))
+    preproc_num_edit=np.mean(preproc_num_dis)
+
+
+
+    '''分段准确率'''
+    test_para_num=np.array(test_para_num)
+    standard_para_num=np.array(standard_para_num)
+    acc_para=np.mean(test_para_num==standard_para_num)
+
+    
+    output=pd.DataFrame()
+    output['总体指标']=[overall_report]
+    output['行内公式平均编辑距离']=[inline_equations_edit]
+    output['行间公式平均编辑距离']=[interline_equations_edit]
+    output['行内公式平均bleu']=[inline_equations_bleu]
+    output['行间公式平均bleu']=[interline_equations_bleu]
+    output['阅读顺序平均编辑距离']=[preproc_num_edit]
+    output['分段准确率']=[acc_para]
+    output['删除的text block的相关指标']=[text_block_report]
+    output['删除的image block的相关指标']=[image_block_report]
+    output['删除的table block的相关指标']=[table_block_report]
+    output['删除的text block的tag相关指标']=[text_block_tag_report]
+    
+
+    return output
+
+"""
+计算编辑距离
+"""
+def Levenshtein_Distance(str1, str2):
+    matrix = [[ i + j for j in range(len(str2) + 1)] for i in range(len(str1) + 1)]
+    for i in range(1, len(str1)+1):
+        for j in range(1, len(str2)+1):
+            if(str1[i-1] == str2[j-1]):
+                d = 0
+            else:
+                d = 1
+            matrix[i][j] = min(matrix[i-1][j]+1, matrix[i][j-1]+1, matrix[i-1][j-1]+d)
+    return matrix[len(str1)][len(str2)]
+
+
+'''
+计算bbox偏移量是否符合标准的函数
+'''
+def bbox_offset(b_t,b_s):
+    '''b_t是test_doc里的bbox,b_s是standard_doc里的bbox'''
+    x1_t,y1_t,x2_t,y2_t=b_t
+    x1_s,y1_s,x2_s,y2_s=b_s
+    x1=max(x1_t,x1_s)
+    x2=min(x2_t,x2_s)
+    y1=max(y1_t,y1_s)
+    y2=min(y2_t,y2_s)
+    area_overlap=(x2-x1)*(y2-y1)
+    area_t=(x2_t-x1_t)*(y2_t-y1_t)+(x2_s-x1_s)*(y2_s-y1_s)-area_overlap
+    if area_t-area_overlap==0 or area_overlap/(area_t-area_overlap)>0.95:
+        return True
+    else:
+        return False
+
+
+
+
+   
+parser = argparse.ArgumentParser()
+parser.add_argument('--test', type=str)
+parser.add_argument('--standard', type=str)
+args = parser.parse_args()
+pdf_json_test = args.test
+pdf_json_standard = args.standard
+
+
+
+if __name__ == '__main__':
+    
+   pdf_json_test = [json.loads(line) 
+                        for line in open(pdf_json_test, 'r', encoding='utf-8')]
+   pdf_json_standard = [json.loads(line) 
+                    for line in open(pdf_json_standard, 'r', encoding='utf-8')]
+   
+   overall_indicator=indicator_cal(pdf_json_standard,pdf_json_test)
+
+   '''计算的指标输出到overall_indicator_output.json中'''
+   overall_indicator.to_json('overall_indicator_output.json',orient='records',lines=True,force_ascii=False)
+