更新scikit-learn依赖兼容性

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 mk_mm_markdown2, mk_nlp_markdown, 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
 
@@ -35,8 +35,8 @@ if __name__ == '__main__':
     # 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)
@@ -58,11 +58,16 @@ if __name__ == '__main__':
             os.makedirs(parent_dir)
 
         # markdown_content = mk_nlp_markdown(pdf_info_dict)
-        markdown_content = mk_mm_markdown2(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:
@@ -55,7 +55,7 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
     return '\n'.join(markdown)
 
 
-def mk_mm_markdown2(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")
@@ -73,16 +73,70 @@ def mk_mm_markdown2(pdf_info_dict:dict):
                     elif span_type == ContentType.InterlineEquation:
                         para_text += f"$$\n{span['content']}\n$$ "
                     elif span_type == ContentType.Image:
-                        para_text += f"![](s3://mllm-raw-media/pdf2md_img/{span['image_path']}) "
+                        para_text += f"![]({join_path(s3_image_save_path, span['image_path'])})"
             markdown.append(para_text)
 
     return '\n\n'.join(markdown)
 
 
-def ocr_mk_mm_standard_format():
+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

@@ -24,6 +24,8 @@ 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/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