更新scikit-learn依赖兼容性

# Conflicts:
#	magic_pdf/dict2md/ocr_mkcontent.py

.gitignore

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

demo/ocr_demo.py

@@ -4,7 +4,7 @@ import os
 from loguru import logger
 from pathlib import Path
 
-from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_nlp_markdown, ocr_mk_mm_markdown
+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
 
@@ -34,8 +34,9 @@ if __name__ == '__main__':
     ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.json"
     # ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.pdf"
     # ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.json"
-    # ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_demo\ocr_1_org.pdf"
-    # ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_demo\ocr_1_fix.json"
+    
+    # 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)
@@ -56,12 +57,17 @@ if __name__ == '__main__':
         if not os.path.exists(parent_dir):
             os.makedirs(parent_dir)
 
-        # markdown_content = ocr_mk_nlp_markdown(pdf_info_dict)
-        markdown_content = ocr_mk_mm_markdown(pdf_info_dict)
+        # 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)
     except Exception as e:

magic_pdf/dict2md/ocr_mkcontent.py

@@ -1,3 +1,4 @@
+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
 
@@ -27,7 +28,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():
@@ -42,7 +42,7 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
                         if not span.get('image_path'):
                             continue
                         else:
-                            content = f"![](s3://mllm-raw-media/pdf2md_img/{span['image_path']})"
+                            content = f"![]({join_path(s3_image_save_path, span['image_path'])})"
                     else:
                         content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                         if span['type'] == ContentType.InlineEquation:
@@ -54,10 +54,89 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
                 markdown.append(line_text.strip() + '  ')
     return '\n'.join(markdown)
 
-def ocr_mk_mm_standard_format():
+
+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:
+                        para_text += span['content']
+                    elif span_type == ContentType.InlineEquation:
+                        para_text += f" ${span['content']}$ "
+                    elif span_type == ContentType.InterlineEquation:
+                        para_text += f"$$\n{span['content']}\n$$ "
+                    elif span_type == ContentType.Image:
+                        para_text += f"![]({join_path(s3_image_save_path, span['image_path'])})"
+            markdown.append(para_text)
+
+    return '\n\n'.join(markdown)
+
+
+def 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:
+                line_text += span['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合并)
-
+    type         string      image/text/table/equation(行间的单独拿出来，行内的和text合并)
+    latex        string      latex文本字段。
+    text         string      纯文本格式的文本数据。
+    md           string      markdown格式的文本数据。
+    img_path     string      s3://full/path/to/img.jpg
     '''
-    pass
+    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/para/para_split.py

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

magic_pdf/pdf_parse_by_model.py

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

@@ -16,6 +16,7 @@ from magic_pdf.libs.commons import (
 from magic_pdf.libs.coordinate_transform import get_scale_ratio
 from magic_pdf.libs.ocr_content_type import ContentType
 from magic_pdf.libs.safe_filename import sanitize_filename
+from magic_pdf.para.para_split import para_split
 from magic_pdf.pre_proc.detect_footer_by_model import parse_footers
 from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model
 from magic_pdf.pre_proc.detect_header import parse_headers
@@ -31,12 +32,13 @@ from magic_pdf.pre_proc.ocr_span_list_modify import remove_spans_by_bboxes, remo
 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,
+def construct_page_component(blocks, para_blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
                              images, tables, interline_equations, inline_equations,
                              dropped_text_block, dropped_image_block, dropped_table_block,
                              need_remove_spans_bboxes_dict):
     return_dict = {
         'preproc_blocks': blocks,
+        "para_blocks": para_blocks, # 分好段落的blocks
         'layout_bboxes': layout_bboxes,
         'page_idx': page_id,
         'page_size': [page_w, page_h],
@@ -234,13 +236,13 @@ def parse_pdf_by_ocr(
         blocks = merge_lines_to_block(lines)
 
         # 根据block合并段落
-
-
+        para_blocks = para_split(blocks, layout_bboxes)
+        
         # 获取QA需要外置的list
         images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
 
         # 构造pdf_info_dict
-        page_info = construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
+        page_info = construct_page_component(blocks, para_blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
                                              images, tables, interline_equations, inline_equations,
                                              dropped_text_block, dropped_image_block, dropped_table_block,
                                              need_remove_spans_bboxes_dict)

magic_pdf/pipeline.py

@@ -3,8 +3,8 @@ 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
+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
@@ -57,7 +57,7 @@ def meta_scan(jso: dict, doc_layout_check=True) -> dict:
     try:
         data_source = get_data_source(jso)
         file_id = jso.get('file_id')
-        book_name = data_source + "/" + file_id
+        book_name = f"{data_source}/{file_id}"
 
         # 首页存在超量drawing问题
         # special_pdf_list = ['zlib/zlib_21822650']
@@ -103,7 +103,7 @@ def classify_by_type(jso: dict, debug_mode=False) -> dict:
         pdf_meta = jso.get('pdf_meta')
         data_source = get_data_source(jso)
         file_id = jso.get('file_id')
-        book_name = data_source + "/" + file_id
+        book_name = f"{data_source}/{file_id}"
         total_page = pdf_meta["total_page"]
         page_width = pdf_meta["page_width_pts"]
         page_height = pdf_meta["page_height_pts"]
@@ -169,7 +169,7 @@ def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:
     try:
         data_source = get_data_source(jso)
         file_id = jso.get('file_id')
-        book_name = data_source + "/" + file_id
+        book_name = f"{data_source}/{file_id}"
         title = jso.get('title')
         url_encode_title = quote(title, safe='')
         if data_source != 'scihub':
@@ -262,7 +262,7 @@ def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
     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
+    book_name = f"{data_source}/{file_id}"
 
     # 1.23.22已修复
     # if debug_mode:
@@ -287,7 +287,7 @@ def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
     #     jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
     else:
         try:
-            save_path = "s3://mllm-raw-media/pdf2md_img/"
+            save_path = s3_image_save_path
             image_s3_config = get_s3_config(save_path)
             start_time = time.time()  # 记录开始时间
             # 先打印一下book_name和解析开始的时间
@@ -326,9 +326,9 @@ def ocr_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
     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
+    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和解析开始的时间
@@ -387,5 +387,31 @@ def ocr_pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
     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
+
+
 if __name__ == "__main__":
     pass

magic_pdf/pre_proc/ocr_dict_merge.py

@@ -24,34 +24,37 @@ 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 = []

magic_pdf/pre_proc/ocr_span_list_modify.py

@@ -77,70 +77,73 @@ 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):

requirements.txt

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