Merge branch 'master' of https://github.com/myhloli/Magic-PDF

demo/demo_test.py

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

demo/ocr_demo.py

@@ -4,8 +4,16 @@ import os
 from loguru import logger
 from pathlib import Path
 
+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.dict2md.ocr_mkcontent import (
+    ocr_mk_mm_markdown_with_para,
+    ocr_mk_nlp_markdown,
+    ocr_mk_mm_markdown,
+    ocr_mk_mm_standard_format,
+    ocr_mk_mm_markdown_with_para_and_pagination,
+    make_standard_format_with_para
+)
 from magic_pdf.libs.commons import join_path
 from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
 
@@ -35,50 +43,59 @@ def ocr_local_parse(ocr_pdf_path, ocr_json_file_path):
         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 = 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))
-
-        # 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):
-    json_object = get_json_from_local_or_s3(book_name)
-    logger.info(json_object)
+    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.get("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)
+    # markdown_pagination = ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict)
+
+    with open(text_content_save_path, "w", encoding="utf-8") as f:
+        f.write(markdown_content)
+
+    standard_format = make_standard_format_with_para(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:
+        # 将standard_format dump成json文本并保存
+        f.write(json.dumps(standard_format, ensure_ascii=False))
+
 
 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/ocr_1.pdf"
-    ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.json"
-    ocr_online_parse(book_name="数学新星网/edu_00001236")
-    ocr_local_parse(ocr_pdf_path, ocr_json_file_path)
+    pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.pdf"
+    json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.json"
+    # ocr_local_parse(pdf_path, json_file_path)
+    book_name = "科数网/edu_00011318"
+    ocr_online_parse(book_name)
+    
     pass

magic_pdf/dict2md/ocr_mkcontent.py

 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):
@@ -58,37 +71,96 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
 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:
+        paras_of_layout = page_info.get("para_blocks")
+        page_markdown = ocr_mk_mm_markdown_with_para_core(paras_of_layout, "mm")
+        markdown.extend(page_markdown)
+    return '\n\n'.join(markdown)
+
+
+def ocr_mk_nlp_markdown_with_para(pdf_info_dict: dict):
+    markdown = []
+    for _, page_info in pdf_info_dict.items():
+        paras_of_layout = page_info.get("para_blocks")
+        page_markdown = ocr_mk_mm_markdown_with_para_core(paras_of_layout, "nlp")
+        markdown.extend(page_markdown)
+    return '\n\n'.join(markdown)
+
+def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: dict):
+    markdown_with_para_and_pagination = []
+    for page_no, page_info in pdf_info_dict.items():
+        paras_of_layout = page_info.get("para_blocks")
+        if not paras_of_layout:
             continue
+        page_markdown = ocr_mk_mm_markdown_with_para_core(paras_of_layout, "mm")
+        markdown_with_para_and_pagination.append({
+            'page_no': page_no,
+            'md_content': '\n\n'.join(page_markdown)
+        })
+    return markdown_with_para_and_pagination
+
+
+def ocr_mk_mm_markdown_with_para_core(paras_of_layout, mode):
+    page_markdown = []
+    for paras in paras_of_layout:
         for para in paras:
             para_text = ''
             for line in para:
                 for span in line['spans']:
                     span_type = span.get('type')
+                    content = ''
                     if span_type == ContentType.Text:
-                        para_text += span['content']
+                        content = ocr_escape_special_markdown_char(split_long_words(span['content']))
                     elif span_type == ContentType.InlineEquation:
-                        para_text += f" ${span['content']}$ "
+                        content = f"${ocr_escape_special_markdown_char(span['content'])}$"
                     elif span_type == ContentType.InterlineEquation:
-                        para_text += f"$$\n{span['content']}\n$$ "
-                    elif span_type == ContentType.Image:
-                        para_text += f"![]({join_path(s3_image_save_path, span['image_path'])})"
-            markdown.append(para_text)
+                        content = f"\n$$\n{ocr_escape_special_markdown_char(span['content'])}\n$$\n"
+                    elif span_type in [ContentType.Image, ContentType.Table]:
+                        if mode == 'mm':
+                            content = f"\n![]({join_path(s3_image_save_path, span['image_path'])})\n"
+                        elif mode == 'nlp':
+                            pass
+                    if content != '':
+                        para_text += content + ' '
+            if para_text.strip() == '':
+                continue
+            else:
+                page_markdown.append(para_text.strip() + '  ')
+    return page_markdown
 
-    return '\n\n'.join(markdown)
 
+def para_to_standard_format(para):
+    para_content = {}
+    if len(para) == 1:
+        para_content = line_to_standard_format(para[0])
+    elif len(para) > 1:
+        para_text = ''
+        inline_equation_num = 0
+        for line in para:
+            for span in line['spans']:
+                span_type = span.get('type')
+                if span_type == ContentType.Text:
+                    content = ocr_escape_special_markdown_char(split_long_words(span['content']))
+                elif span_type == ContentType.InlineEquation:
+                    content = f"${ocr_escape_special_markdown_char(span['content'])}$"
+                    inline_equation_num += 1
+                para_text += content + ' '
+        para_content = {
+            'type': 'text',
+            'text': para_text,
+            'inline_equation_num': inline_equation_num
+        }
+    return para_content
 
 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:
+        paras_of_layout = page_info.get("para_blocks")
+        if not paras_of_layout:
             continue
-        for para in paras:
-            for line in para:
-                content = line_to_standard_format(line)
-                content_list.append(content)
+        for paras in paras_of_layout:
+            for para in paras:
+                para_content = para_to_standard_format(para)
+                content_list.append(para_content)
     return content_list
 
 
@@ -125,7 +197,8 @@ def line_to_standard_format(line):
                 line_text += f"${inline_equation}$"
                 inline_equation_num += 1
             elif span['type'] == ContentType.Text:
-                line_text += span['content']
+                text_content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
+                line_text += text_content
     content = {
         'type': 'text',
         'text': line_text,

magic_pdf/libs/boxbase.py

@@ -18,6 +18,33 @@ def _is_in_or_part_overlap(box1, box2) -> bool:
                 y1_1 < y0_2 or  # box1在box2的上边
                 y0_1 > y1_2)    # box1在box2的下边
 
+def _is_in_or_part_overlap_with_area_ratio(box1, box2, area_ratio_threshold=0.6):
+    """
+    判断box1是否在box2里面，或者box1和box2有部分重叠，且重叠面积占box1的比例超过area_ratio_threshold
+    
+    """
+    if box1 is None or box2 is None:
+        return False
+    
+    x0_1, y0_1, x1_1, y1_1 = box1
+    x0_2, y0_2, x1_2, y1_2 = box2
+
+    if not _is_in_or_part_overlap(box1, box2):
+        return False
+    
+    # 计算重叠面积
+    x_left = max(x0_1, x0_2)
+    y_top = max(y0_1, y0_2)
+    x_right = min(x1_1, x1_2)
+    y_bottom = min(y1_1, y1_2)
+    overlap_area = (x_right - x_left) * (y_bottom - y_top)
+    
+    # 计算box1的面积
+    box1_area = (x1_1 - x0_1) * (y1_1 - y0_1)
+    
+    return overlap_area / box1_area > area_ratio_threshold
+    
+    
 def _is_in(box1, box2) -> bool:
     """
     box1是否完全在box2里面

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/pdf_parse_by_ocr.py

@@ -57,16 +57,16 @@ def construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, lay
 
 
 def parse_pdf_by_ocr(
-    pdf_path,
-    s3_pdf_profile,
-    pdf_model_output,
-    save_path,
-    book_name,
-    pdf_model_profile=None,
-    image_s3_config=None,
-    start_page_id=0,
-    end_page_id=None,
-    debug_mode=False,
+        pdf_path,
+        s3_pdf_profile,
+        pdf_model_output,
+        save_path,
+        book_name,
+        pdf_model_profile=None,
+        image_s3_config=None,
+        start_page_id=0,
+        end_page_id=None,
+        debug_mode=False,
 ):
     pdf_bytes = read_file(pdf_path, s3_pdf_profile)
     save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
@@ -95,7 +95,6 @@ def parse_pdf_by_ocr(
 
     start_time = time.time()
 
-
     end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
     for page_id in range(start_page_id, end_page_id + 1):
 
@@ -125,13 +124,6 @@ def parse_pdf_by_ocr(
             page_id, page, ocr_page_info, md_bookname_save_path, debug_mode=debug_mode
         )
 
-        # 构建需要remove的bbox列表
-        # need_remove_spans_bboxes = []
-        # need_remove_spans_bboxes.extend(page_no_bboxes)
-        # need_remove_spans_bboxes.extend(header_bboxes)
-        # need_remove_spans_bboxes.extend(footer_bboxes)
-        # need_remove_spans_bboxes.extend(footnote_bboxes)
-
         # 构建需要remove的bbox字典
         need_remove_spans_bboxes_dict = {
             DropTag.PAGE_NUMBER: page_no_bboxes,
@@ -199,50 +191,48 @@ def parse_pdf_by_ocr(
             else:
                 continue
 
-
-
-
-        # 删除重叠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相关数据
+        '''
+        删除remove_span_block_bboxes中的bbox
+        并增加drop相关数据
+        '''
         spans, dropped_spans_by_removed_bboxes = remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict)
 
-        # 对image和table截图
+        '''对image和table截图'''
         spans = cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client)
 
-        # 行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)
+        '''行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)'''
         displayed_list = []
         text_inline_lines = []
         modify_y_axis(spans, displayed_list, text_inline_lines)
-        # 模型识别错误的行间公式, type类型转换成行内公式
+
+        '''模型识别错误的行间公式, type类型转换成行内公式'''
         spans = modify_inline_equation(spans, displayed_list, text_inline_lines)
 
-        # bbox去除粘连
+        '''bbox去除粘连'''
         spans = remove_overlap_between_bbox(spans)
 
-        # 对tpye=["interline_equation", "image", "table"]进行额外处理,如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
+        '''
+        对tpye=["interline_equation", "image", "table"]进行额外处理,
+        如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
+        '''
         spans = adjust_bbox_for_standalone_block(spans)
 
-
-        # 从ocr_page_info中解析layout信息(按自然阅读方向排序,并修复重叠和交错的bad case)
+        '''从ocr_page_info中解析layout信息(按自然阅读方向排序,并修复重叠和交错的bad case)'''
         layout_bboxes, layout_tree = layout_detect(ocr_page_info['subfield_dets'], page, ocr_page_info)
 
-        # 将spans合并成line(在layout内,从上到下,从左到右)
+        '''将spans合并成line(在layout内,从上到下,从左到右)'''
         lines, dropped_spans_by_layout = merge_spans_to_line_by_layout(spans, layout_bboxes)
 
-        # 将lines合并成block
+        '''将lines合并成block'''
         blocks = merge_lines_to_block(lines)
 
-        # 根据block合并段落
-        #para_blocks = para_split(blocks, layout_bboxes)
-        
-        # 获取QA需要外置的list
+        '''获取QA需要外置的list'''
         images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
 
-        # drop的span_list合并
+        '''drop的span_list合并'''
         dropped_spans = []
         dropped_spans.extend(dropped_spans_by_span_overlap)
         dropped_spans.extend(dropped_spans_by_removed_bboxes)
@@ -263,19 +253,18 @@ def parse_pdf_by_ocr(
             elif span['type'] in [ContentType.InlineEquation, ContentType.InterlineEquation]:
                 dropped_equation_block.append(span)
 
-
-
-        # 构造pdf_info_dict
+        '''构造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_equation_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)
-    
-    # 在测试时,保存调试信息
+    para_split(pdf_info_dict, debug_mode=debug_mode)
+
+    '''在测试时,保存调试信息'''
     if debug_mode:
         params_file_save_path = join_path(
             save_tmp_path, "md", book_name, "preproc_out.json"
@@ -284,7 +273,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/pipeline_txt.py

+"""
+文本型pdf转化为统一清洗格式
+"""
+
+
+
+from loguru import logger
+from magic_pdf.dict2md.mkcontent import mk_universal_format
+from magic_pdf.libs.commons import join_path
+from magic_pdf.libs.json_compressor import JsonCompressor
+from magic_pdf.spark.base import exception_handler, get_data_source
+
+
+def txt_pdf_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")
+            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 = mk_universal_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)}",)
+        # 把无用的信息清空
+        jso["doc_layout_result"] = ""
+        jso["pdf_intermediate_dict"] = ""
+        jso["pdf_meta"] = ""
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso

magic_pdf/pre_proc/ocr_dict_merge.py

@@ -66,7 +66,7 @@ def merge_spans_to_line_by_layout(spans, layout_bboxes):
         # 遍历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:

magic_pdf/pre_proc/ocr_span_list_modify.py

@@ -44,10 +44,15 @@ def remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict):
         # logger.info(f"remove spans by bbox dict, drop_tag: {drop_tag}, removed_bboxes: {removed_bboxes}")
         need_remove_spans = []
         for span in spans:
+            # 通过判断span的bbox是否在removed_bboxes中, 判断是否需要删除该span
             for removed_bbox in removed_bboxes:
                 if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], removed_bbox) > 0.5:
                     need_remove_spans.append(span)
                     break
+                # 当drop_tag为DropTag.FOOTNOTE时, 判断span是否在removed_bboxes中任意一个的下方，如果是,则删除该span
+                elif drop_tag == DropTag.FOOTNOTE and (span['bbox'][1]+span['bbox'][3])/2 > removed_bbox[3] and removed_bbox[0] < (span['bbox'][0]+span['bbox'][2])/2 < removed_bbox[2]:
+                    need_remove_spans.append(span)
+                    break
 
         for span in need_remove_spans:
             spans.remove(span)

magic_pdf/spark/base.py

+
+
+from loguru import logger
+
+from magic_pdf.libs.drop_reason import DropReason
+
+
+def get_data_source(jso: dict):
+    data_source = jso.get("data_source")
+    if data_source is None:
+        data_source = jso.get("file_source")
+    return data_source
+
+
+def exception_handler(jso: dict, e):
+    logger.exception(e)
+    jso["need_drop"] = True
+    jso["drop_reason"] = DropReason.Exception
+    jso["exception"] = f"ERROR: {e}"
+    return jso
+

magic_pdf/train_utils/convert_to_train_format.py

+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"]:
+            if "paras" in para:
+                paras = para["paras"]
+                for para_key, para_content in paras.items():
+                    para_bbox = para_content["para_bbox"]
+                    is_para_title = para_content["is_para_title"]
+                    if is_para_title:
+                        n_bbox = {"category_id": 0, "bbox": para_bbox}
+                    else:
+                        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)
+
+        for footnote_bbox in v["bak_footer_note_bboxes"]:
+            n_bbox = {"category_id": 5, "bbox": list(footnote_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

+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

+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

+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

@@ -2,7 +2,7 @@ boto3>=1.28.43
 Brotli>=1.1.0
 click>=8.1.7
 Distance>=0.1.3
-PyMuPDF>=1.23.26
+PyMuPDF>=1.24.0
 loguru>=0.6.0
 matplotlib>=3.8.3
 numpy>=1.21.6
@@ -12,5 +12,6 @@ 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
\ No newline at end of file