增加uni_parse_pdf逻辑

# Conflicts:
#	magic_pdf/libs/drop_tag.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,
@@ -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

@@ -30,13 +30,13 @@ def read_json_file(file_path):
 
 
 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\双栏\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/j.1540-627x.2006.00176.x.pdf"
+    ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.json"
     try:
         ocr_pdf_model_info = read_json_file(ocr_json_file_path)
         pth = Path(ocr_json_file_path)

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

@@ -3,11 +3,12 @@ import numpy as np
 from loguru import logger
 
 from magic_pdf.libs.boxbase import _is_in
+from magic_pdf.libs.ocr_content_type import ContentType
 
 
 LINE_STOP_FLAG = ['.', '!', '?', '。', '！', '？',"：", ":", ")", "）", ";"]
-INLINE_EQUATION = 'inline_equation'
-INTER_EQUATION = "displayed_equation"
+INLINE_EQUATION = ContentType.InlineEquation
+INTERLINE_EQUATION = ContentType.InterlineEquation
 TEXT = "text"
 
 def __add_line_period(blocks, layout_bboxes):
@@ -20,20 +21,19 @@ def __add_line_period(blocks, layout_bboxes):
         for line in block['lines']:
             last_span = line['spans'][-1]
             span_type = last_span['type']
-            if span_type in [TEXT, INLINE_EQUATION]:
+            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, INTER_EQUATION]:
+                    if span_type in [INLINE_EQUATION, INTERLINE_EQUATION]:
                         last_span['content'] = span_content + '.'
 
 
 
 def __valign_lines(blocks, layout_bboxes):
     """
-    对齐行的左侧和右侧。
-    扫描行的左侧和右侧，如果x0, x1差距不超过3就强行对齐到所处layout的左右两侧（和layout有一段距离）。
-    3是个经验值，TODO，计算得来
-    
+    在一个layoutbox内对齐行的左侧和右侧。
+    扫描行的左侧和右侧，如果x0, x1差距不超过一个阈值，就强行对齐到所处layout的左右两侧（和layout有一段距离）。
+    3是个经验值，TODO，计算得来，可以设置为1.5个正文字符。
     """
     
     min_distance = 3
@@ -159,11 +159,14 @@ def __split_para_in_layoutbox(lines_group, layout_bboxes, lang="en", char_avg_le
                 else: 
                     para.append(line)
             else: # 其他，图片、表格、行间公式，各自占一段
-                para.append(line)
-                paras.append(para)
+                if len(para)>0:
+                    paras.append(para)
+                    para = []
+                else:
+                    paras.append([line])
+                    para = []
                 # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
                 # logger.info(para_text)
-                para = []
         if len(para)>0:
             paras.append(para)
             # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])

magic_pdf/pdf_parse_by_ocr.py

@@ -14,6 +14,7 @@ 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
@@ -34,7 +35,7 @@ from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
 
 def construct_page_component(blocks, para_blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
                              images, tables, interline_equations, inline_equations,
-                             dropped_text_block, dropped_image_block, dropped_table_block,
+                             dropped_text_block, dropped_image_block, dropped_table_block, dropped_equation_block,
                              need_remove_spans_bboxes_dict):
     return_dict = {
         'preproc_blocks': blocks,
@@ -50,6 +51,7 @@ def construct_page_component(blocks, para_blocks, layout_bboxes, page_id, page_w
         '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
@@ -133,10 +135,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"]
@@ -202,12 +204,12 @@ def parse_pdf_by_ocr(
 
 
         # 删除重叠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, img_s3_client)
@@ -230,7 +232,7 @@ 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)
@@ -241,10 +243,33 @@ def parse_pdf_by_ocr(
         # 获取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, para_blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
                                              images, tables, interline_equations, inline_equations,
-                                             dropped_text_block, dropped_image_block, dropped_table_block,
+                                             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
 

magic_pdf/pipeline.py

@@ -7,7 +7,7 @@ from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_nlp_markdown, ocr_mk_mm_markd
 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
@@ -237,9 +237,10 @@ def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> 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"] = ""
@@ -413,5 +414,65 @@ def ocr_pdf_intermediate_dict_to_standard_format(jso: dict, debug_mode=False) ->
     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
+
+
+def ocr_dropped_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
+    # 检测debug开关
+    if debug_mode:
+        pass
+    else:  # 如果debug没开，则检测是否有needdrop字段
+        if not jso.get('need_drop', False):
+            return jso
+        else:
+            s3_pdf_path = jso.get('file_location')
+            s3_config = get_s3_config(s3_pdf_path)
+            model_output_json_list = jso.get('doc_layout_result')
+            data_source = get_data_source(jso)
+            file_id = jso.get('file_id')
+            book_name = f"{data_source}/{file_id}"
+            try:
+                save_path = s3_image_save_path
+                image_s3_config = get_s3_config(save_path)
+                start_time = time.time()  # 记录开始时间
+                # 先打印一下book_name和解析开始的时间
+                logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
+                pdf_info_dict = parse_pdf_by_ocr(
+                    s3_pdf_path,
+                    s3_config,
+                    model_output_json_list,
+                    save_path,
+                    book_name,
+                    pdf_model_profile=None,
+                    image_s3_config=image_s3_config,
+                    start_page_id=start_page_id,
+                    debug_mode=debug_mode
+                )
+                if pdf_info_dict.get('need_drop', False):  # 如果返回的字典里有need_drop，则提取drop_reason并跳过本次解析
+                    jso['need_drop'] = True
+                    jso['drop_reason'] = pdf_info_dict["drop_reason"]
+                else:  # 正常返回，将 pdf_info_dict 压缩并存储
+                    pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
+                    jso['pdf_intermediate_dict'] = pdf_info_dict
+                end_time = time.time()  # 记录完成时间
+                parse_time = int(end_time - start_time)  # 计算执行时间
+                # 解析完成后打印一下book_name和耗时
+                logger.info(
+                    f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
+                    file=sys.stderr)
+                jso['parse_time'] = parse_time
+            except Exception as e:
+                jso = exception_handler(jso, e)
+                return jso
+
+
 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_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
 
 
@@ -59,6 +60,7 @@ def merge_spans_to_line(spans):
 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中
@@ -78,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):
@@ -98,7 +95,8 @@ def modify_y_axis(spans: list, displayed_list: list, text_inline_lines: list):
             # 如果当前的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):
+                    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)
@@ -160,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:

requirements.txt

@@ -11,6 +11,6 @@ pycld2>=0.41
 regex>=2023.12.25
 spacy>=3.7.4
 termcolor>=2.4.0
-scikit-learn>=1.4.1.post1
+scikit-learn
 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