增加ocr版本解析功能

.github/workflows/python-package.yml

@@ -20,21 +20,29 @@ jobs:
         python-version: ["3.10"]
 
     steps:
-    - uses: actions/checkout@v4
+    - name: Checkout code
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+
     - name: Install wheel
       run: |
         python -m pip install wheel
+
     - name: Build wheel
       run: |
         python setup.py bdist_wheel
+
     - name: Upload artifact
       uses: actions/upload-artifact@v4
       with:

README.md

@@ -19,6 +19,17 @@ python 3.9+
 git clone https://github.com/myhloli/Magic-PDF.git
 ```
 
+2.Install the requirements
+
+```sh
+pip install -r requirements.txt
+```
+
+3.Run the main script
+
+```sh
+use demo/demo_test.py
+```
 
 ### 版权说明
 

demo/download.py

@@ -2,7 +2,7 @@ import json
 import os
 from tqdm import tqdm
 
-from magic_pdf.libs import join_path
+from magic_pdf.libs.commons import join_path
 
 with open('/mnt/petrelfs/share_data/ouyanglinke/OCR/OCR_validation_dataset.json', 'r') as f:
     samples = json.load(f)

demo/draw_bbox.py

@@ -1,4 +1,4 @@
-from magic_pdf.libs import fitz  # PyMuPDF
+from magic_pdf.libs.commons import fitz  # PyMuPDF
 
 # PDF文件路径
 pdf_path = "D:\\project\\20231108code-clean\\code-clean\\tmp\\unittest\\download-pdfs\\scihub\\scihub_53700000\\libgen.scimag53724000-53724999.zip_10.1097\\00129191-200509000-00018.pdf"

demo/ocr_demo.py

@@ -0,0 +1,29 @@
+import os
+
+from loguru import logger
+
+from magic_pdf.dict2md.ocr_mkcontent import mk_nlp_markdown
+from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
+
+
+def save_markdown(markdown_text, input_filepath):
+    # 获取输入文件的目录
+    directory = os.path.dirname(input_filepath)
+    # 获取输入文件的文件名（不带扩展名）
+    base_name = os.path.basename(input_filepath)
+    file_name_without_ext = os.path.splitext(base_name)[0]
+    # 定义输出文件的路径
+    output_filepath = os.path.join(directory, f"{file_name_without_ext}.md")
+
+    # 将Markdown文本写入.md文件
+    with open(output_filepath, 'w', encoding='utf-8') as file:
+        file.write(markdown_text)
+
+
+if __name__ == '__main__':
+    ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_0.json"
+    pdf_info_dict = parse_pdf_by_ocr(ocr_json_file_path)
+    markdown_text = mk_nlp_markdown(pdf_info_dict)
+    logger.info(markdown_text)
+    save_markdown(markdown_text, ocr_json_file_path)
+

demo/pdf2md.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import click
 from loguru import logger
 
-from magic_pdf.libs import join_path
+from magic_pdf.libs.commons import join_path
 from magic_pdf.dict2md.mkcontent import mk_mm_markdown
 from magic_pdf.pipeline import parse_pdf_by_model
 

magic_pdf/dict2md/ocr_mkcontent.py

@@ -0,0 +1,21 @@
+def mk_nlp_markdown(pdf_info_dict: dict):
+
+    markdown = []
+
+    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']:
+                line_text = ''
+                for span in line['spans']:
+                    content = span['content'].replace('$', '\$')  # 转义$
+                    if span['type'] == 'inline_equation':
+                        content = f"${content}$"
+                    elif span['type'] == 'displayed_equation':
+                        content = f"$$\n{content}\n$$"
+                    line_text += content + ' '
+                # 在行末添加两个空格以强制换行
+                markdown.append(line_text.strip() + '  ')
+    return '\n'.join(markdown)

magic_pdf/libs/boxbase.py

@@ -119,6 +119,20 @@ def _is_left_overlap(box1, box2,):
     return x0_1<=x0_2<=x1_1 and vertical_overlap_cond
 
 
+def __is_overlaps_y_exceeds_threshold(bbox1, bbox2, overlap_ratio_threshold=0.8):
+    """检查两个bbox在y轴上是否有重叠，并且该重叠区域的高度占两个bbox高度更低的那个超过80%"""
+    _, y0_1, _, y1_1 = bbox1
+    _, y0_2, _, y1_2 = bbox2
+
+    overlap = max(0, min(y1_1, y1_2) - max(y0_1, y0_2))
+    height1, height2 = y1_1 - y0_1, y1_2 - y0_2
+    max_height = max(height1, height2)
+    min_height = min(height1, height2)
+
+    return (overlap / min_height) > overlap_ratio_threshold
+
+
+
 def calculate_iou(bbox1, bbox2):
     # Determine the coordinates of the intersection rectangle
     x_left = max(bbox1[0], bbox2[0])
@@ -163,7 +177,25 @@ def calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2):
     else:
         return intersection_area / min_box_area
 
-    
+
+def get_minbox_if_overlap_by_ratio(bbox1, bbox2, ratio):
+    """
+    通过calculate_overlap_area_2_minbox_area_ratio计算两个bbox重叠的面积占最小面积的box的比例
+    如果比例大于ratio，则返回小的那个bbox,
+    否则返回None
+    """
+    x1_min, y1_min, x1_max, y1_max = bbox1
+    x2_min, y2_min, x2_max, y2_max = bbox2
+    area1 = (x1_max - x1_min) * (y1_max - y1_min)
+    area2 = (x2_max - x2_min) * (y2_max - y2_min)
+    overlap_ratio = calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
+    if overlap_ratio > ratio and area1 < area2:
+        return bbox1
+    elif overlap_ratio > ratio and area2 < area1:
+        return bbox2
+    else:
+        return None
+
 def get_bbox_in_boundry(bboxes:list, boundry:tuple)-> list:
     x0, y0, x1, y1 = boundry
     new_boxes = [box for box in bboxes if box[0] >= x0 and box[1] >= y0 and box[2] <= x1 and box[3] <= y1]

magic_pdf/libs/ocr_dict_merge.py

@@ -0,0 +1,46 @@
+from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold
+
+
+def merge_spans(spans):
+    # 按照y0坐标排序
+    spans.sort(key=lambda span: span['bbox'][1])
+
+    lines = []
+    current_line = [spans[0]]
+    for span in spans[1:]:
+        # 如果当前的span类型为"displayed_equation" 或者 当前行中已经有"displayed_equation"
+        if span['type'] == "displayed_equation" or any(s['type'] == "displayed_equation" 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)
+
+    # 计算每行的边界框，并对每行中的span按照x0进行排序
+    line_objects = []
+    for line in lines:
+        # 按照x0坐标排序
+        line.sort(key=lambda span: span['bbox'][0])
+        line_bbox = [
+            min(span['bbox'][0] for span in line),  # x0
+            min(span['bbox'][1] for span in line),  # y0
+            max(span['bbox'][2] for span in line),  # x1
+            max(span['bbox'][3] for span in line),  # y1
+        ]
+        line_objects.append({
+            "bbox": line_bbox,
+            "spans": line,
+        })
+
+    return line_objects

magic_pdf/pdf_parse_by_ocr.py

@@ -0,0 +1,85 @@
+import json
+
+from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio
+from magic_pdf.libs.ocr_dict_merge import merge_spans
+
+
+def read_json_file(file_path):
+    with open(file_path, 'r') as f:
+        data = json.load(f)
+    return data
+
+
+def construct_page_component(page_id, text_blocks_preproc):
+    return_dict = {
+        'preproc_blocks': text_blocks_preproc,
+        'page_idx': page_id
+    }
+    return return_dict
+
+
+def parse_pdf_by_ocr(
+    ocr_json_file_path,
+    start_page_id=0,
+    end_page_id=None,
+):
+    ocr_pdf_info = read_json_file(ocr_json_file_path)
+    pdf_info_dict = {}
+    end_page_id = end_page_id if end_page_id else len(ocr_pdf_info) - 1
+    for page_id in range(start_page_id, end_page_id + 1):
+        ocr_page_info = ocr_pdf_info[page_id]
+        layout_dets = ocr_page_info['layout_dets']
+        spans = []
+        for layout_det in layout_dets:
+            category_id = layout_det['category_id']
+            allow_category_id_list = [13, 14, 15]
+            if category_id in allow_category_id_list:
+                x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
+                bbox = [int(x0), int(y0), int(x1), int(y1)]
+                #  13: 'embedding',     # 嵌入公式
+                #  14: 'isolated',      # 单行公式
+                #  15: 'ocr_text',      # ocr识别文本
+                span = {
+                    'bbox': bbox,
+                }
+                if category_id == 13:
+                    span['content'] = layout_det['latex']
+                    span['type'] = 'inline_equation'
+                elif category_id == 14:
+                    span['content'] = layout_det['latex']
+                    span['type'] = 'displayed_equation'
+                elif category_id == 15:
+                    span['content'] = layout_det['text']
+                    span['type'] = 'text'
+                # print(span)
+                spans.append(span)
+            else:
+                continue
+
+        # 合并重叠的spans
+        for span1 in spans.copy():
+            for span2 in spans.copy():
+                if span1 != span2:
+                    overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.8)
+                    if overlap_box is not None:
+                        bbox_to_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
+                        if bbox_to_remove is not None:
+                            spans.remove(bbox_to_remove)
+
+        # 将spans合并成line
+        lines = merge_spans(spans)
+
+        # 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
+        blocks = []
+        for line in lines:
+            blocks.append({
+                "bbox": line['bbox'],
+                "lines": [line],
+            })
+
+        # 构造pdf_info_dict
+        page_info = construct_page_component(page_id, blocks)
+        pdf_info_dict[f"page_{page_id}"] = page_info
+
+    return pdf_info_dict
+

setup.py

@@ -1,5 +1,5 @@
 from setuptools import setup, find_packages
-
+import subprocess
 def parse_requirements(filename):
     with open(filename) as f:
         lines = f.read().splitlines()
@@ -15,12 +15,26 @@ def parse_requirements(filename):
 
     return requires
 
+def get_version():
+    command = ["git", "describe", "--tags"]
+    try:
+        version = subprocess.check_output(command).decode().strip()
+        version_parts = version.split("-")
+        if len(version_parts) > 1 and version_parts[0].startswith("magic_pdf"):
+            return version_parts[1]
+        else:
+            raise ValueError(f"Invalid version tag {version}. Expected format is magic_pdf-<version>-released.")
+    except Exception as e:
+        print(e)
+        return "0.0.0"
+
 
 requires = parse_requirements('requirements.txt')
 
 setup(
     name="magic_pdf",  # 项目名
-    version="0.1.3",  # 版本号
+    # version="0.1.3",  # 版本号
+    version=get_version(),  # 自动从tag中获取版本号
     packages=find_packages(),  # 包含所有的包
     install_requires=requires,  # 项目依赖的第三方库
     python_requires=">=3.9",  # 项目依赖的 Python 版本

tests/test_commons.py

@@ -1,10 +1,10 @@
 import io
 import json
 import os
-from magic_pdf.libs import fitz
+from magic_pdf.libs.commons import fitz
 
 from app.common.s3 import get_s3_config, get_s3_client
-from magic_pdf.libs import join_path, json_dump_path, read_file, parse_bucket_key
+from magic_pdf.libs.commons import join_path, json_dump_path, read_file, parse_bucket_key
 from loguru import logger
 
 test_pdf_dir_path = "s3://llm-pdf-text/unittest/pdf/"

tests/test_metascan_classify/test_classify.py

@@ -2,10 +2,10 @@ import os
 
 import pytest
 
-from magic_pdf.filter import classify_by_area, classify_by_text_len, classify_by_avg_words, \
+from magic_pdf.filter.pdf_classify_by_type import classify_by_area, classify_by_text_len, classify_by_avg_words, \
     classify_by_img_num, classify_by_text_layout, classify_by_img_narrow_strips
 from magic_pdf.filter.pdf_meta_scan import get_pdf_page_size_pts, get_pdf_textlen_per_page, get_imgs_per_page
-from test.test_commons import get_docs_from_test_pdf, get_test_json_data
+from tests.test_commons import get_docs_from_test_pdf, get_test_json_data
 
 # 获取当前目录
 current_directory = os.path.dirname(os.path.abspath(__file__))

tests/test_metascan_classify/test_meta_scan.py

@@ -2,7 +2,7 @@ import os
 
 import pytest
 from magic_pdf.filter.pdf_meta_scan import get_pdf_page_size_pts, get_image_info, get_pdf_text_layout_per_page, get_language
-from test.test_commons import get_docs_from_test_pdf, get_test_json_data
+from tests.test_commons import get_docs_from_test_pdf, get_test_json_data
 
 # 获取当前目录
 current_directory = os.path.dirname(os.path.abspath(__file__))

tests/test_para/test_para_pipeline.py

@@ -11,21 +11,21 @@ Execute the following command to run the tests under directory code-clean:
     
 """
 
-from test.test_para.test_pdf2text_recogPara_Common import (
+from tests.test_para.test_pdf2text_recogPara_Common import (
     TestIsBboxOverlap,
     TestIsInBbox,
     TestIsBboxOverlap,
     TestIsLineLeftAlignedFromNeighbors,
     TestIsLineRightAlignedFromNeighbors,
 )
-from test.test_para.test_pdf2text_recogPara_EquationsProcessor import TestCalcOverlapPct
-from test.test_para.test_pdf2text_recogPara_BlockInnerParasProcessor import TestIsConsistentLines
-from test.test_para.test_pdf2text_recogPara_BlockContinuationProcessor import (
+from tests.test_para.test_pdf2text_recogPara_EquationsProcessor import TestCalcOverlapPct
+from tests.test_para.test_pdf2text_recogPara_BlockInnerParasProcessor import TestIsConsistentLines
+from tests.test_para.test_pdf2text_recogPara_BlockContinuationProcessor import (
     TestIsAlphabetChar,
     TestIsChineseChar,
     TestIsOtherLetterChar,
 )
-from test.test_para.test_pdf2text_recogPara_TitleProcessor import TestTitleProcessor
+from tests.test_para.test_pdf2text_recogPara_TitleProcessor import TestTitleProcessor
 
 
 # Test suite