diff --git a/.gitignore b/.gitignore index 6de528ed..d90d0c33 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,5 @@ ocr_demo /app/common/__init__.py /magic_pdf/config/__init__.py +source.dev.env + diff --git a/README_zh-CN.md b/README_zh-CN.md index fdab67bd..708c140f 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -121,7 +121,9 @@ pip install magic-pdf[full]==0.6.2b1 -i https://pypi.tuna.tsinghua.edu.cn/simple #### 2. 下载模型权重文件 详细参考 [如何下载模型文件](docs/how_to_download_models_zh_cn.md) -下载后请将models目录移动到空间较大的ssd磁盘目录 +> ❗️模型下载后请务必检查模型文件是否下载完整 +> +> 请检查目录下的模型文件大小与网页上描述是否一致，如果可以的话，最好通过sha256校验模型是否下载完整 #### 3. 拷贝配置文件并进行配置 在仓库根目录可以获得 [magic-pdf.template.json](magic-pdf.template.json) 配置模版文件 @@ -133,9 +135,11 @@ cp magic-pdf.template.json ~/magic-pdf.json ``` 在用户目录中找到magic-pdf.json文件并配置"models-dir"为[2. 下载模型权重文件](#2-下载模型权重文件)中下载的模型权重文件所在目录 -> ❗️务必正确配置模型权重文件所在目录，否则会因为找不到模型文件而导致程序无法运行 +> ❗️务必正确配置模型权重文件所在目录的【绝对路径】，否则会因为找不到模型文件而导致程序无法运行 > -> windows系统中应把路径中所有的"\\"替换为"/",否则会因为转义原因导致json文件语法错误。 +> windows系统中此路径应包含盘符，且需把路径中所有的"\"替换为"/",否则会因为转义原因导致json文件语法错误。 +> +> 例如：模型放在D盘根目录的models目录，则model-dir的值应为"D:/models" ```json { "models-dir": "/tmp/models" diff --git a/README_zh-CN_v2.md b/README_zh-CN_v2.md index 755a2321..3c4688af 100644 --- a/README_zh-CN_v2.md +++ b/README_zh-CN_v2.md @@ -16,12 +16,11 @@ opendatalab%2FMinerU | Trendshift -[English](README.md) | [简体中文](README_zh-CN.md) | [日本語](README_ja-JP.md) +[English](README.md) | [简体中文](README_zh-CN.md)

-MinerU: 端到端的PDF解析工具（基于PDF-Extract-Kit）支持PDF转Markdown🚀🚀🚀
PDF-Extract-Kit: 高质量PDF解析工具箱🔥🔥🔥

@@ -35,7 +34,8 @@ # 更新记录 -- 2024/07/18 首次开源 +- 2024/08/01 0.6.2b1发布，优化了依赖冲突问题和安装文档 +- 2024/07/05 首次开源 @@ -82,7 +82,10 @@ # MinerU ## 项目简介 MinerU是一款将PDF转化为机器可读格式的工具（如markdown、json），可以很方便地抽取为任意格式。 -MinerU诞生于[书生-浦语](https://github.com/InternLM/InternLM)的预训练过程中，我们将会集中精力解决科技文献中的符号转化问题，以此在大模型时代为科技发展做出一点贡献。 +MinerU诞生于[书生-浦语](https://github.com/InternLM/InternLM)的预训练过程中，我们将会集中精力解决科技文献中的符号转化问题，希望在大模型时代为科技发展做出贡献。 +相比国内外知名商用产品MinerU还很年轻，如果遇到问题或者结果不及预期请到issue提交问题，同时附上相关PDF。 + +https://github.com/user-attachments/assets/4bea02c9-6d54-4cd6-97ed-dff14340982c ## 主要功能 @@ -102,12 +105,12 @@ MinerU诞生于[书生-浦语](https://github.com/InternLM/InternLM)的预训练 如果遇到任何安装问题，请先查询 FAQ
如果遇到解析效果不及预期，参考 Known Issue
有3种不同方式可以体验MinerU的效果： -- 在线体验(无需任何安装) -- 使用CPU快速体验（Windows，Linux，Mac） -- Linux/Windows + CUDA +- [在线体验(无需任何安装)](#在线体验) +- [使用CPU快速体验（Windows，Linux，Mac）](#使用cpu快速体验) +- [Linux/Windows + CUDA](#使用gpu) -**软硬件环境支持说明** +**⚠️安装前必看——软硬件环境支持说明** 为了确保项目的稳定性和可靠性，我们在开发过程中仅对特定的软硬件环境进行优化和测试。这样当用户在推荐的系统配置上部署和运行项目时，能够获得最佳的性能表现和最少的兼容性问题。 @@ -166,23 +169,47 @@ MinerU诞生于[书生-浦语](https://github.com/InternLM/InternLM)的预训练 ### 在线体验 -[在线体验点击这里](TODO) +[在线体验点击这里](https://opendatalab.com/OpenSourceTools/Extractor/PDF) ### 使用CPU快速体验 +#### 1. 安装magic-pdf ```bash pip install magic-pdf[full]==0.6.2b1 detectron2 --extra-index-url https://wheels.myhloli.com -i https://pypi.tuna.tsinghua.edu.cn/simple ``` +#### 2. 下载模型权重文件 -> ❗️已收到多起由于镜像源和依赖冲突问题导致安装了错误版本软件包的反馈，请务必安装完成后通过以下命令验证版本是否正确 -> ```bash -> magic-pdf --version -> ``` -> 如版本低于0.6.2b1，请提交issue进行反馈。 +详细参考 [如何下载模型文件](docs/how_to_download_models_zh_cn.md) +> ❗️模型下载后请务必检查模型文件是否下载完整 +> +> 请检查目录下的模型文件大小与网页上描述是否一致，如果可以的话，最好通过sha256校验模型是否下载完整 + +#### 3. 拷贝配置文件并进行配置 +在仓库根目录可以获得 [magic-pdf.template.json](magic-pdf.template.json) 配置模版文件 +> ❗️务必执行以下命令将配置文件拷贝到【用户目录】下，否则程序将无法运行 +> +> windows的用户目录为 "C:\Users\用户名", linux用户目录为 "/home/用户名", macOS用户目录为 "/Users/用户名" +```bash +cp magic-pdf.template.json ~/magic-pdf.json +``` + +在用户目录中找到magic-pdf.json文件并配置"models-dir"为[2. 下载模型权重文件](#2-下载模型权重文件)中下载的模型权重文件所在目录 +> ❗️务必正确配置模型权重文件所在目录的【绝对路径】，否则会因为找不到模型文件而导致程序无法运行 +> +> windows系统中此路径应包含盘符，且需把路径中所有的"\"替换为"/",否则会因为转义原因导致json文件语法错误。 +> +> 例如：模型放在D盘根目录的models目录，则model-dir的值应为"D:/models" +```json +{ + "models-dir": "/tmp/models" +} +``` ### 使用GPU +如果您的设备支持CUDA，且满足主线环境中的显卡要求，则可以使用GPU加速，请根据自己的系统选择适合的教程： + - [Ubuntu22.04LTS + GPU](docs/README_Ubuntu_CUDA_Acceleration_zh_CN.md) - [Windows10/11 + GPU](docs/README_Windows_CUDA_Acceleration_zh_CN.md) @@ -191,7 +218,45 @@ pip install magic-pdf[full]==0.6.2b1 detectron2 --extra-index-url https://wheels ### 命令行 -TODO +```bash +magic-pdf --help +Usage: magic-pdf [OPTIONS] + +Options: + -v, --version display the version and exit + -p, --path PATH local pdf filepath or directory [required] + -o, --output-dir TEXT output local directory + -m, --method [ocr|txt|auto] the method for parsing pdf. + ocr: using ocr technique to extract information from pdf, + txt: suitable for the text-based pdf only and outperform ocr, + auto: automatically choose the best method for parsing pdf + from ocr and txt. + without method specified, auto will be used by default. + --help Show this message and exit. + + +## show version +magic-pdf -v + +## command line example +magic-pdf -p {some_pdf} -o {some_output_dir} -m auto +``` + +其中 `{some_pdf}` 可以是单个pdf文件，也可以是一个包含多个pdf文件的目录。 +运行完命令后输出的结果会保存在`{some_output_dir}`目录下, 输出的文件列表如下 + +```text +├── some_pdf.md # markdown 文件 +├── images # 存放图片目录 +├── layout.pdf # layout 绘图 +├── middle.json # minerU 中间处理结果 +├── model.json # 模型推理结果 +├── origin.pdf # 原 pdf 文件 +└── spans.pdf # 最小粒度的bbox位置信息绘图 +``` + +更多有关输出文件的信息，请参考[输出文件说明](docs/output_file_zh_cn.md) + ### API @@ -245,6 +310,8 @@ TODO - 阅读顺序基于规则的分割，在一些情况下会乱序 - 列表、代码块、目录在layout模型里还没有支持 - 漫画书、艺术图册、小学教材、习题尚不能很好解析 +- 在一些公式密集的PDF上强制启用OCR效果会更好 +- 如果您要处理包含大量公式的pdf,强烈建议开启OCR功能。使用pymuPDF提取文字的时候会出现文本行互相重叠的情况导致公式插入位置不准确。 好消息是，这些我们正在努力实现！ @@ -267,6 +334,7 @@ The project currently leverages PyMuPDF to deliver advanced functionalities; how # Acknowledgments +- [StructEqTable](https://github.com/UniModal4Reasoning/StructEqTable-Deploy) - [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR) - [PyMuPDF](https://github.com/pymupdf/PyMuPDF) - [fast-langdetect](https://github.com/LlmKira/fast-langdetect) diff --git a/demo/small_ocr.pdf b/demo/small_ocr.pdf new file mode 100644 index 00000000..2ab92332 Binary files /dev/null and b/demo/small_ocr.pdf differ diff --git a/docs/FAQ_zh_cn.md b/docs/FAQ_zh_cn.md index b5932b60..48fda7b6 100644 --- a/docs/FAQ_zh_cn.md +++ b/docs/FAQ_zh_cn.md @@ -50,20 +50,7 @@ pip install paddlepaddle==3.0.0b1 可能是由于模型文件未下载完整导致，可尝试重现下载模型文件后再试 参考：https://github.com/opendatalab/MinerU/issues/143 -### 7.程序运行完成后，找不到tmp目录 - -程序输出目录是在"magic-pdf.json"中通过 -```json -{ - "temp-output-dir": "/tmp" -} -``` -进行配置的。 -如果没有更改这个参数，使用默认的配置执行程序，在linux/macOS会在绝对路径"/tmp"下创建一个"magic-pdf"文件夹作为输出路径。 -而在windows下，默认的输出路径与执行命令时，命令行所在的盘符相关，如果命令行在C盘，则默认输出路径为"C://tmp/magic-pdf"。 -参考：https://github.com/opendatalab/MinerU/issues/149 - -### 8.模型文件应该下载到哪里/models-dir的配置应该怎么填 +### 7.模型文件应该下载到哪里/models-dir的配置应该怎么填 模型文件的路径输入是在"magic-pdf.json"中通过 ```json @@ -75,24 +62,16 @@ pip install paddlepaddle==3.0.0b1 这个路径是绝对路径而不是相对路径，绝对路径的获取可在models目录中通过命令 "pwd" 获取。 参考：https://github.com/opendatalab/MinerU/issues/155#issuecomment-2230216874 -### 9.命令行中 --model "model_json_path" 指的是什么？ +### 8.命令行中 --model "model_json_path" 指的是什么？ model_json 指的是通过模型分析后生成的一种有特定格式的json文件。 如果使用 https://github.com/opendatalab/PDF-Extract-Kit 项目生成，该文件一般在项目的output目录下。 如果使用 MinerU 的命令行调用内置的模型分析，该文件一般在输出路径"/tmp/magic-pdf/pdf-name"下。 参考：https://github.com/opendatalab/MinerU/issues/128 -### 10.Linux下报错：Required dependency not installed, please install by "pip install magic-pdf[full-cpu] detectron2 --extra-index-url https://myhloli.github.io/wheels/" +### 9.报错：Required dependency not installed, please install by "pip install magic-pdf[full-cpu] detectron2 --extra-index-url https://myhloli.github.io/wheels/" -这种情况可以先使用pip list 检查一下自己的依赖库列表，重点确认下以下几个库有没有安装（版本不一定完全一致，有就可以） +通过更新0.6.2b1来解决 ```bash -opencv-contrib-python 4.6.0.66 -opencv-python 4.6.0.66 -opencv-python-headless 4.10.0.84 -paddleocr 2.7.3 -paddlepaddle 2.6.1 -torch 2.2.2 -torchtext 0.17.2 -torchvision 0.17.2 -``` -如果都有的话，可能是libgl库没有安装，参考 https://github.com/opendatalab/MinerU/issues/165#issuecomment-2245202282 安装libgl库后再试试能不能正常使用。 +pip install magic-pdf[full]==0.6.2b1 -i https://pypi.tuna.tsinghua.edu.cn/simple +``` \ No newline at end of file diff --git a/docs/README_Ubuntu_CUDA_Acceleration_zh_CN.md b/docs/README_Ubuntu_CUDA_Acceleration_zh_CN.md index e69de29b..4f5634ee 100644 --- a/docs/README_Ubuntu_CUDA_Acceleration_zh_CN.md +++ b/docs/README_Ubuntu_CUDA_Acceleration_zh_CN.md @@ -0,0 +1,114 @@ +# Ubuntu 22.04 LTS + +## 1. 更新apt +```bash +sudo apt-get update +``` +## 2. 检测是否已安装nvidia驱动 +```bash +nvidia-smi +``` +如果看到类似如下的信息，说明已经安装了nvidia驱动，可以跳过步骤3 +``` ++---------------------------------------------------------------------------------------+ +| NVIDIA-SMI 537.34 Driver Version: 537.34 CUDA Version: 12.2 | +|-----------------------------------------+----------------------+----------------------+ +| GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC | +| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | +| | | MIG M. | +|=========================================+======================+======================| +| 0 NVIDIA GeForce RTX 3060 Ti WDDM | 00000000:01:00.0 On | N/A | +| 0% 51C P8 12W / 200W | 1489MiB / 8192MiB | 5% Default | +| | | N/A | ++-----------------------------------------+----------------------+----------------------+ +``` +## 3. 安装驱动 +如没有驱动，则通过如下命令 +```bash +sudo apt-get install nvidia-driver-545 +``` +安装专有驱动，安装完成后，重启电脑 +```bash +reboot +``` +## 4. 安装anacoda +如果已安装conda，可以跳过本步骤 +```bash +wget https://repo.anaconda.com/archive/Anaconda3-2024.06-1-Linux-x86_64.sh +bash Anaconda3-2024.06-1-Linux-x86_64.sh +``` +最后一步输入yes，关闭终端重新打开 +## 5. 使用conda 创建环境 +需指定python版本为3.10 +```bash +conda create -n MinerU python=3.10 +conda activate MinerU +``` +## 6. 安装应用 +```bash +pip install magic-pdf[full]==0.6.2b1 detectron2 --extra-index-url https://wheels.myhloli.com -i https://pypi.tuna.tsinghua.edu.cn/simple +``` +> ❗️下载完成后，务必通过以下命令确认magic-pdf的版本是否正确 +> +> ```bash +> magic-pdf --version +>``` +> 如果版本号小于0.6.2，请到issue中向我们反馈 + +## 7. 下载模型 +详细参考 [如何下载模型文件](how_to_download_models_zh_cn.md) +下载后请将models目录移动到空间较大的ssd磁盘目录 +> ❗️模型下载后请务必检查模型文件是否下载完整 +> +> 请检查目录下的模型文件大小与网页上描述是否一致，如果可以的话，最好通过sha256校验模型是否下载完整 +> +## 8. 第一次运行前的配置 +在仓库根目录可以获得 [magic-pdf.template.json](../magic-pdf.template.json) 配置模版文件 +> ❗️务必执行以下命令将配置文件拷贝到【用户目录】下，否则程序将无法运行 +> +> linux用户目录为 "/home/用户名" +```bash +wget https://github.com/opendatalab/MinerU/raw/master/magic-pdf.template.json +cp magic-pdf.template.json ~/magic-pdf.json +``` + +在用户目录中找到magic-pdf.json文件并配置"models-dir"为[7. 下载模型](#7-下载模型)中下载的模型权重文件所在目录 +> ❗️务必正确配置模型权重文件所在目录的【绝对路径】，否则会因为找不到模型文件而导致程序无法运行 +> +```json +{ + "models-dir": "/tmp/models" +} +``` + +## 9. 第一次运行 +从仓库中下载样本文件，并测试 +```bash +wget https://github.com/opendatalab/MinerU/raw/master/demo/small_ocr.pdf +magic-pdf pdf-command --pdf small_ocr.pdf +``` +## 10. 测试CUDA加速 +如果您的显卡显存大于等于8G，可以进行以下流程，测试CUDA解析加速效果 + +**1.修改【用户目录】中配置文件magic-pdf.json中"device-mode"的值** +```json +{ + "device-mode":"cuda" +} +``` +**2.运行以下命令测试cuda加速效果** +```bash +magic-pdf pdf-command --pdf small_ocr.pdf +``` + +## 11. 为ocr开启cuda加速 +> ❗️以下操作需显卡显存大于等于16G才可进行，否则会因为显存不足导致程序崩溃或运行速度下降 + +**1.下载paddlepaddle-gpu, 安装完成后会自动开启ocr加速** +```bash +python -m pip install paddlepaddle-gpu==3.0.0b1 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/ +``` +**2.运行以下命令测试ocr加速效果** +```bash +magic-pdf pdf-command --pdf small_ocr.pdf +``` diff --git a/docs/README_Windows_CUDA_Acceleration_zh_CN.md b/docs/README_Windows_CUDA_Acceleration_zh_CN.md index e69de29b..ad43c83a 100644 --- a/docs/README_Windows_CUDA_Acceleration_zh_CN.md +++ b/docs/README_Windows_CUDA_Acceleration_zh_CN.md @@ -0,0 +1,103 @@ +# Windows10/11 + +## 1. 安装cuda和cuDNN + +需要安装的版本 CUDA 11.8 + cuDNN 8.7.0 +- CUDA 11.8 https://developer.nvidia.com/cuda-11-8-0-download-archive +- cuDNN v8.7.0 (November 28th, 2022), for CUDA 11.x https://developer.nvidia.com/rdp/cudnn-archive + +## 2. 安装anaconda +如果已安装conda，可以跳过本步骤 + +下载链接： +https://repo.anaconda.com/archive/Anaconda3-2024.06-1-Windows-x86_64.exe + +## 3. 使用conda 创建环境 +需指定python版本为3.10 +```bash +conda create -n MinerU python=3.10 +conda activate MinerU +``` +## 4. 安装应用 +```bash +pip install magic-pdf[full]==0.6.2b1 detectron2 --extra-index-url https://wheels.myhloli.com -i https://pypi.tuna.tsinghua.edu.cn/simple +``` +> ❗️下载完成后，务必通过以下命令确认magic-pdf的版本是否正确 +> +> ```bash +> magic-pdf --version +>``` +> 如果版本号小于0.6.2，请到issue中向我们反馈 + +## 5. 下载模型 +详细参考 [如何下载模型文件](how_to_download_models_zh_cn.md) +下载后请将models目录移动到空间较大的ssd磁盘目录 +> ❗️模型下载后请务必检查模型文件是否下载完整 +> +> 请检查目录下的模型文件大小与网页上描述是否一致，如果可以的话，最好通过sha256校验模型是否下载完整 + +## 6. 第一次运行前的配置 +在仓库根目录可以获得 [magic-pdf.template.json](../magic-pdf.template.json) 配置模版文件 +> ❗️务必执行以下命令将配置文件拷贝到【用户目录】下，否则程序将无法运行 +> +> windows用户目录为 "C:\Users\用户名" +```powershell +(New-Object System.Net.WebClient).DownloadFile('https://github.com/opendatalab/MinerU/raw/master/magic-pdf.template.json', 'magic-pdf.template.json') +cp magic-pdf.template.json ~/magic-pdf.json +``` + +在用户目录中找到magic-pdf.json文件并配置"models-dir"为[5. 下载模型](#5-下载模型)中下载的模型权重文件所在目录 +> ❗️务必正确配置模型权重文件所在目录的【绝对路径】，否则会因为找不到模型文件而导致程序无法运行 +> +> windows系统中此路径应包含盘符，且需把路径中所有的"\"替换为"/",否则会因为转义原因导致json文件语法错误。 +> +> 例如：模型放在D盘根目录的models目录，则model-dir的值应为"D:/models" +```json +{ + "models-dir": "/tmp/models" +} +``` + +## 7. 第一次运行 +从仓库中下载样本文件，并测试 +```powershell +(New-Object System.Net.WebClient).DownloadFile('https://github.com/opendatalab/MinerU/raw/master/demo/small_ocr.pdf', 'small_ocr.pdf') +magic-pdf pdf-command --pdf small_ocr.pdf +``` + +## 8. 测试CUDA加速 +如果您的显卡显存大于等于8G，可以进行以下流程，测试CUDA解析加速效果 + +**1.覆盖安装支持cuda的torch和torchvision** +```bash +pip install --force-reinstall torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118 +``` +> ❗️务必在命令中指定以下版本 +> ```bash +> torch==2.3.1 torchvision==0.18.1 +> ``` +> 这是我们支持的最高版本，如果不指定版本会自动安装更高版本导致程序无法运行 + +**2.修改【用户目录】中配置文件magic-pdf.json中"device-mode"的值** +```json +{ + "device-mode":"cuda" +} +``` +**3.运行以下命令测试cuda加速效果** +```bash +magic-pdf pdf-command --pdf small_ocr.pdf +``` + +## 9. 为ocr开启cuda加速 +> ❗️以下操作需显卡显存大于等于16G才可进行，否则会因为显存不足导致程序崩溃或运行速度下降 + +**1.下载paddlepaddle-gpu, 安装完成后会自动开启ocr加速** +```bash +pip install paddlepaddle-gpu==2.6.1 +``` +**2.运行以下命令测试ocr加速效果** +```bash +magic-pdf pdf-command --pdf small_ocr.pdf +``` + diff --git a/docs/how_to_download_models_en.md b/docs/how_to_download_models_en.md index 7ae49824..80a85bb5 100644 --- a/docs/how_to_download_models_en.md +++ b/docs/how_to_download_models_en.md @@ -35,7 +35,7 @@ model_dir = snapshot_download('wanderkid/PDF-Extract-Kit') Alternatively, you can use Git to clone the model repository from ModelScope: ```bash -git clone https://www.modelscope.cn/wanderkid/PDF-Extract-Kit.git +git lfs clone https://www.modelscope.cn/wanderkid/PDF-Extract-Kit.git ``` @@ -45,7 +45,7 @@ Put [model files]() here: ./ ├── Layout │ ├── config.json -│ └── weights.pth +│ └── model_final.pth ├── MFD │ └── weights.pt ├── MFR @@ -57,4 +57,4 @@ Put [model files]() here: │ ├── tokenizer_config.json │ └── tokenizer.json └── README.md -``` \ No newline at end of file +``` diff --git a/docs/how_to_download_models_zh_cn.md b/docs/how_to_download_models_zh_cn.md index 464124c4..c37716c5 100644 --- a/docs/how_to_download_models_zh_cn.md +++ b/docs/how_to_download_models_zh_cn.md @@ -1,23 +1,37 @@ -### 安装 Git LFS -开始之前，请确保您的系统上已安装 Git 大文件存储 (Git LFS)。使用以下命令进行安装 +# 如何下载模型文件 + +模型文件可以从Hugging Face 或 Model Scope 下载，由于网络原因，国内用户访问HF 可能会失败，请使用 ModelScope。 + + +方法一：[从 Hugging Face 下载模型](#方法一从-hugging-face-下载模型) + +方法二：[从 ModelScope 下载模型](#方法二从-modelscope-下载模型) + +## 方法一：从 Hugging Face 下载模型 + +使用Git LFS 从Hugging Face下载模型文件 + +```bash +git lfs install # 安装 Git 大文件存储插件 (Git LFS) +git lfs clone https://huggingface.co/wanderkid/PDF-Extract-Kit # 从 Hugging Face 下载 PDF-Extract-Kit 模型 +``` + + +## 方法二：从 ModelScope 下载模型 +ModelScope 支持SDK或模型下载，任选一个即可。 + +[Git lsf下载](#1利用git-lsf下载) + +[SDK下载](#2利用sdk下载) + +### 1）利用Git lsf下载 ```bash git lfs install +git lfs clone https://www.modelscope.cn/wanderkid/PDF-Extract-Kit.git ``` -### 从 Hugging Face 下载模型 -请使用以下命令从 Hugging Face 下载 PDF-Extract-Kit 模型： - -```bash -git lfs clone https://huggingface.co/wanderkid/PDF-Extract-Kit -``` - -确保在克隆过程中启用了 Git LFS，以便正确下载所有大文件。 - - -### 从 ModelScope 下载模型 - -#### SDK下载 +### 2）利用SDK下载 ```bash # 首先安装modelscope @@ -30,42 +44,18 @@ from modelscope import snapshot_download model_dir = snapshot_download('wanderkid/PDF-Extract-Kit') ``` -#### Git下载 -也可以使用git clone从 ModelScope 下载模型: - -需要先安装git lfs - ->##### On Linux -> ->Debian and RPM packages are available from packagecloud, see the [Linux installation instructions](INSTALLING.md). -> ->##### On macOS -> ->[Homebrew](https://brew.sh) bottles are distributed and can be installed via `brew install git-lfs`. -> ->##### On Windows -> ->Git LFS is included in the distribution of [Git for Windows](https://gitforwindows.org/). ->Alternatively, you can install a recent version of Git LFS from the [Chocolatey](https://chocolatey.org/) package manager. - -然后通过git clone下载模型: -```bash -git lfs clone https://www.modelscope.cn/wanderkid/PDF-Extract-Kit.git -``` - - -将 'models' 目录移动到具有较大磁盘空间的目录中，最好是在固态硬盘(SSD)上。 - +## 额外步骤 +### 1.检查模型目录是否下载完整 模型文件夹的结构如下，包含了不同组件的配置文件和权重文件： ``` ./ -├── Layout +├── Layout # 布局检测模型 │ ├── config.json │ └── model_final.pth -├── MFD +├── MFD # 公式检测 │ └── weights.pt -├── MFR +├── MFR # 公式识别模型 │ └── UniMERNet │ ├── config.json │ ├── preprocessor_config.json @@ -73,7 +63,7 @@ git lfs clone https://www.modelscope.cn/wanderkid/PDF-Extract-Kit.git │ ├── README.md │ ├── tokenizer_config.json │ └── tokenizer.json -│── TabRec +│── TabRec # 表格识别模型 │ └─StructEqTable │ ├── config.json │ ├── generation_config.json @@ -85,3 +75,10 @@ git lfs clone https://www.modelscope.cn/wanderkid/PDF-Extract-Kit.git │ └── tokenizer_config.json └── README.md ``` + +### 2.检查模型文件是否下载完整 +请检查目录下的模型文件大小与网页上描述是否一致，如果可以的话，最好通过sha256校验模型是否下载完整 + +### 3.移动模型到固态硬盘 +将 'models' 目录移动到具有较大磁盘空间的目录中，最好是在固态硬盘(SSD)上。 +此外在 `~/magic-pdf.json`里修改模型的目录指向最终的模型存放位置，否则会报模型无法加载的错误。 diff --git a/docs/images/layout_example.png b/docs/images/layout_example.png new file mode 100644 index 00000000..14e2116f Binary files /dev/null and b/docs/images/layout_example.png differ diff --git a/docs/images/poly.png b/docs/images/poly.png new file mode 100644 index 00000000..14af7726 Binary files /dev/null and b/docs/images/poly.png differ diff --git a/docs/images/spans_example.png b/docs/images/spans_example.png new file mode 100644 index 00000000..14de87ed Binary files /dev/null and b/docs/images/spans_example.png differ diff --git a/docs/output_file_zh_cn.md b/docs/output_file_zh_cn.md new file mode 100644 index 00000000..393e8e0d --- /dev/null +++ b/docs/output_file_zh_cn.md @@ -0,0 +1,332 @@ + + +## 概览 +`magic-pdf` 命令执行后除了输出和 markdown 有关的文件以外，还会生成若干个和 markdown 无关的文件。现在将一一介绍这些文件 + + +### layout.pdf +每一页的 layout 均由一个或多个框组成。 每个框左上脚的数字表明它们的序号。此外 layout.pdf 框内用不同的背景色块圈定不同的内容块。 + +![layout 页面示例](images/layout_example.png) + + +### spans.pdf +根据 span 类型的不同，采用不同颜色线框绘制页面上所有 span。该文件可以用于质检，可以快速排查出文本丢失、行间公式未识别等问题。 + +![span 页面示例](images/spans_example.png) + + +### model.json + +#### 结构定义 +```python +from pydantic import BaseModel, Field +from enum import IntEnum + +class CategoryType(IntEnum): + title = 0 # 标题 + plain_text = 1 # 文本 + abandon = 2 # 包括页眉页脚页码和页面注释 + figure = 3 # 图片 + figure_caption = 4 # 图片描述 + table = 5 # 表格 + table_caption = 6 # 表格描述 + table_footnote = 7 # 表格注释 + isolate_formula = 8 # 行间公式 + formula_caption = 9 # 行间公式的标号 + + embedding = 13 # 行内公式 + isolated = 14 # 行间公式 + text = 15 # ocr 识别结果 + + +class PageInfo(BaseModel): + page_no: int = Field(description="页码序号，第一页的序号是 0", ge=0) + height: int = Field(description="页面高度", gt=0) + width: int = Field(description="页面宽度", ge=0) + +class ObjectInferenceResult(BaseModel): + category_id: CategoryType = Field(description="类别", ge=0) + poly: list[float] = Field(description="四边形坐标, 分别是 左上，右上，右下，左下 四点的坐标") + score: float = Field(description="推理结果的置信度") + latex: str | None = Field(description="latex 解析结果", default=None) + html: str | None = Field(description="html 解析结果", default=None) + +class PageInferenceResults(BaseModel): + layout_dets: list[ObjectInferenceResult] = Field(description="页面识别结果", ge=0) + page_info: PageInfo = Field(description="页面元信息") + + +# 所有页面的推理结果按照页码顺序依次放到列表中即为 minerU 推理结果 +inference_result: list[PageInferenceResults] = [] + +``` + +poly 坐标的格式 [x0, y0, x1, y1, x2, y2, x3, y3], 分别表示左上、右上、右下、左下四点的坐标 +![poly 坐标示意图](images/poly.png) + + +#### 示例数据 + +```json +[ + { + "layout_dets": [ + { + "category_id": 2, + "poly": [ + 99.1906967163086, + 100.3119125366211, + 730.3707885742188, + 100.3119125366211, + 730.3707885742188, + 245.81326293945312, + 99.1906967163086, + 245.81326293945312 + ], + "score": 0.9999997615814209 + } + ], + "page_info": { + "page_no": 0, + "height": 2339, + "width": 1654 + } + }, + { + "layout_dets": [ + { + "category_id": 5, + "poly": [ + 99.13092803955078, + 2210.680419921875, + 497.3183898925781, + 2210.680419921875, + 497.3183898925781, + 2264.78076171875, + 99.13092803955078, + 2264.78076171875 + ], + "score": 0.9999997019767761 + } + ], + "page_info": { + "page_no": 1, + "height": 2339, + "width": 1654 + } + } +] +``` + + +### middle.json + +| 字段名 | 解释 | +| :-----|:------------------------------------------| +|pdf_info | list，每个元素都是一个dict,这个dict是每一页pdf的解析结果，详见下表 | +|_parse_type | ocr \| txt，用来标识本次解析的中间态使用的模式 | +|_version_name | string, 表示本次解析使用的 magic-pdf 的版本号 | + +
+ +**pdf_info** +字段结构说明 + +| 字段名 | 解释 | +| :-----| :---- | +| preproc_blocks | pdf预处理后，未分段的中间结果 | +| layout_bboxes | 布局分割的结果，含有布局的方向（垂直、水平），和bbox，按阅读顺序排序 | +| page_idx | 页码，从0开始 | +| page_size | 页面的宽度和高度 | +| _layout_tree | 布局树状结构 | +| images | list，每个元素是一个dict，每个dict表示一个img_block | +| tables | list，每个元素是一个dict，每个dict表示一个table_block | +| interline_equations | list，每个元素是一个dict，每个dict表示一个interline_equation_block | +| discarded_blocks | List, 模型返回的需要drop的block信息 | +| para_blocks | 将preproc_blocks进行分段之后的结果 | + +上表中 `para_blocks` 是个dict的数组，每个dict是一个block结构，block最多支持一次嵌套 + +
+ +**block** + +外层block被称为一级block，一级block中的字段包括 + +| 字段名 | 解释 | +| :-----| :---- | +| type | block类型（table\|image）| +|bbox | block矩形框坐标 | +|blocks |list，里面的每个元素都是一个dict格式的二级block | + +
+一级block只有"table"和"image"两种类型，其余block均为二级block + +二级block中的字段包括 + +| 字段名 | 解释 | +| :-----| :---- | +| type | block类型 | +| bbox | block矩形框坐标 | +| lines | list，每个元素都是一个dict表示的line，用来描述一行信息的构成| + +二级block的类型详解 + +| type | desc | +|:-------------------| :---- | +| image_body | 图像的本体 | +| image_caption | 图像的描述文本 | +| table_body | 表格本体 | +| table_caption | 表格的描述文本 | +| table_footnote | 表格的脚注 | +| text | 文本块 | +| title | 标题块 | +| interline_equation | 行间公式块| + +
+ +**line** + +line 的 字段格式如下 + +| 字段名 | 解释 | +| :-----| :---- | +| bbox | line的矩形框坐标 | +| spans | list，每个元素都是一个dict表示的span，用来描述一个最小组成单元的构成 | + + +
+ +**span** + +| 字段名 | 解释 | +| :-----| :---- | +| bbox | span的矩形框坐标 | +| type | span的类型 | +| content \| img_path | 文本类型的span使用content，图表类使用img_path 用来存储实际的文本或者截图路径信息 | + +span 的类型有如下几种 + +| type | desc | +| :-----| :---- | +| image | 图片 | +| table | 表格 | +| text | 文本 | +| inline_equation | 行内公式 | +| interline_equation | 行间公式 | + + +**总结** + +span是所有元素的最小存储单元 + +para_blocks内存储的元素为区块信息 + +区块结构为 + +一级block(如有)->二级block->line->span + + +#### 示例数据 + +```json +{ + "pdf_info": [ + { + "preproc_blocks": [ + { + "type": "text", + "bbox": [ + 52, + 61.956024169921875, + 294, + 82.99800872802734 + ], + "lines": [ + { + "bbox": [ + 52, + 61.956024169921875, + 294, + 72.0000228881836 + ], + "spans": [ + { + "bbox": [ + 54.0, + 61.956024169921875, + 296.2261657714844, + 72.0000228881836 + ], + "content": "dependent on the service headway and the reliability of the departure ", + "type": "text", + "score": 1.0 + } + ] + } + ] + } + ], + "layout_bboxes": [ + { + "layout_bbox": [ + 52, + 61, + 294, + 731 + ], + "layout_label": "V", + "sub_layout": [] + } + ], + "page_idx": 0, + "page_size": [ + 612.0, + 792.0 + ], + "_layout_tree": [], + "images": [], + "tables": [], + "interline_equations": [], + "discarded_blocks": [], + "para_blocks": [ + { + "type": "text", + "bbox": [ + 52, + 61.956024169921875, + 294, + 82.99800872802734 + ], + "lines": [ + { + "bbox": [ + 52, + 61.956024169921875, + 294, + 72.0000228881836 + ], + "spans": [ + { + "bbox": [ + 54.0, + 61.956024169921875, + 296.2261657714844, + 72.0000228881836 + ], + "content": "dependent on the service headway and the reliability of the departure ", + "type": "text", + "score": 1.0 + } + ] + } + ] + } + ] + } + ], + "_parse_type": "txt", + "_version_name": "0.6.1" +} +``` \ No newline at end of file diff --git a/magic-pdf.template.json b/magic-pdf.template.json index dacb6d92..f6adfbfa 100644 --- a/magic-pdf.template.json +++ b/magic-pdf.template.json @@ -3,7 +3,6 @@ "bucket-name-1":["ak", "sk", "endpoint"], "bucket-name-2":["ak", "sk", "endpoint"] }, - "temp-output-dir":"/tmp", "models-dir":"/tmp/models", "device-mode":"cpu", "table-config": { diff --git a/magic_pdf/cli/magicpdf.py b/magic_pdf/cli/magicpdf.py deleted file mode 100644 index 45dc5eb1..00000000 --- a/magic_pdf/cli/magicpdf.py +++ /dev/null @@ -1,359 +0,0 @@ -""" -这里实现2个click命令： -第一个： - 接收一个完整的s3路径，例如：s3://llm-pdf-text/pdf_ebook_and_paper/pre-clean-mm-markdown/v014/part-660420b490be-000008.jsonl?bytes=0,81350 - 1）根据~/magic-pdf.json里的ak,sk等，构造s3cliReader读取到这个jsonl的对应行，返回json对象。 - 2）根据Json对象里的pdf的s3路径获取到他的ak,sk,endpoint，构造出s3cliReader用来读取pdf - 3）从magic-pdf.json里读取到本地保存图片、Md等的临时目录位置,构造出LocalImageWriter，用来保存截图 - 4）从magic-pdf.json里读取到本地保存图片、Md等的临时目录位置,构造出LocalIRdWriter，用来读写本地文件 - - 最后把以上步骤准备好的对象传入真正的解析API - -第二个： - 接收1）pdf的本地路径。2）模型json文件（可选）。然后： - 1）根据~/magic-pdf.json读取到本地保存图片、md等临时目录的位置，构造出LocalImageWriter，用来保存截图 - 2）从magic-pdf.json里读取到本地保存图片、Md等的临时目录位置,构造出LocalIRdWriter，用来读写本地文件 - 3）根据约定，根据pdf本地路径，推导出pdf模型的json，并读入 - - -效果： -python magicpdf.py json-command --json s3://llm-pdf-text/scihub/xxxx.json?bytes=0,81350 -python magicpdf.py pdf-command --pdf /home/llm/Downloads/xxxx.pdf --model /home/llm/Downloads/xxxx.json 或者 python magicpdf.py --pdf /home/llm/Downloads/xxxx.pdf -""" - -import os -import json as json_parse -import click -from loguru import logger -from pathlib import Path -from magic_pdf.libs.version import __version__ - -from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode -from magic_pdf.libs.draw_bbox import draw_layout_bbox, draw_span_bbox -from magic_pdf.pipe.UNIPipe import UNIPipe -from magic_pdf.pipe.OCRPipe import OCRPipe -from magic_pdf.pipe.TXTPipe import TXTPipe -from magic_pdf.libs.path_utils import ( - parse_s3path, - parse_s3_range_params, - remove_non_official_s3_args, -) -from magic_pdf.libs.config_reader import ( - get_local_dir, - get_s3_config, -) -from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter -from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter -from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter -import csv -import copy -import magic_pdf.model as model_config - -parse_pdf_methods = click.Choice(["ocr", "txt", "auto"]) - - -def prepare_env(pdf_file_name, method): - local_parent_dir = os.path.join(get_local_dir(), "magic-pdf", pdf_file_name, method) - - local_image_dir = os.path.join(str(local_parent_dir), "images") - local_md_dir = local_parent_dir - os.makedirs(local_image_dir, exist_ok=True) - os.makedirs(local_md_dir, exist_ok=True) - return local_image_dir, local_md_dir - - -def write_to_csv(csv_file_path, csv_data): - with open(csv_file_path, mode="a", newline="", encoding="utf-8") as csvfile: - # 创建csv writer对象 - csv_writer = csv.writer(csvfile) - # 写入数据 - csv_writer.writerow(csv_data) - logger.info(f"数据已成功追加到 '{csv_file_path}'") - - -def do_parse( - pdf_file_name, - pdf_bytes, - model_list, - parse_method, - f_draw_span_bbox=True, - f_draw_layout_bbox=True, - f_dump_md=True, - f_dump_middle_json=True, - f_dump_model_json=True, - f_dump_orig_pdf=True, - f_dump_content_list=True, - f_make_md_mode=MakeMode.MM_MD, -): - - orig_model_list = copy.deepcopy(model_list) - - local_image_dir, local_md_dir = prepare_env(pdf_file_name, parse_method) - image_writer, md_writer = DiskReaderWriter(local_image_dir), DiskReaderWriter(local_md_dir) - image_dir = str(os.path.basename(local_image_dir)) - - if parse_method == "auto": - jso_useful_key = {"_pdf_type": "", "model_list": model_list} - pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True) - elif parse_method == "txt": - pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True) - elif parse_method == "ocr": - pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True) - else: - logger.error("unknown parse method") - exit(1) - - pipe.pipe_classify() - - """如果没有传入有效的模型数据，则使用内置model解析""" - if len(model_list) == 0: - if model_config.__use_inside_model__: - pipe.pipe_analyze() - orig_model_list = copy.deepcopy(pipe.model_list) - else: - logger.error("need model list input") - exit(1) - - pipe.pipe_parse() - pdf_info = pipe.pdf_mid_data["pdf_info"] - if f_draw_layout_bbox: - draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir) - if f_draw_span_bbox: - draw_span_bbox(pdf_info, pdf_bytes, local_md_dir) - - md_content = pipe.pipe_mk_markdown(image_dir, drop_mode=DropMode.NONE, md_make_mode=f_make_md_mode) - if f_dump_md: - """写markdown""" - md_writer.write( - content=md_content, - path=f"{pdf_file_name}.md", - mode=AbsReaderWriter.MODE_TXT, - ) - - if f_dump_middle_json: - """写middle_json""" - md_writer.write( - content=json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4), - path=f"{pdf_file_name}_middle.json", - mode=AbsReaderWriter.MODE_TXT, - ) - - if f_dump_model_json: - """写model_json""" - md_writer.write( - content=json_parse.dumps(orig_model_list, ensure_ascii=False, indent=4), - path=f"{pdf_file_name}_model.json", - mode=AbsReaderWriter.MODE_TXT, - ) - - if f_dump_orig_pdf: - """写源pdf""" - md_writer.write( - content=pdf_bytes, - path=f"{pdf_file_name}_origin.pdf", - mode=AbsReaderWriter.MODE_BIN, - ) - - content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE) - if f_dump_content_list: - """写content_list""" - md_writer.write( - content=json_parse.dumps(content_list, ensure_ascii=False, indent=4), - path=f"{pdf_file_name}_content_list.json", - mode=AbsReaderWriter.MODE_TXT, - ) - logger.info(f"local output dir is '{local_md_dir}', you can found the result in it.") - - -@click.group() -@click.version_option(__version__, "--version", "-v", help="显示版本信息") -@click.help_option("--help", "-h", help="显示帮助信息") -def cli(): - pass - - -@cli.command() -@click.option("--json", type=str, help="输入一个S3路径") -@click.option( - "--method", - type=parse_pdf_methods, - help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法", - default="auto", -) -@click.option("--inside_model", type=click.BOOL, default=True, help="使用内置模型测试") -@click.option("--model_mode", type=click.STRING, default="full", - help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢") -def json_command(json, method, inside_model, model_mode): - model_config.__use_inside_model__ = inside_model - model_config.__model_mode__ = model_mode - - if not json.startswith("s3://"): - logger.error("usage: magic-pdf json-command --json s3://some_bucket/some_path") - exit(1) - - def read_s3_path(s3path): - bucket, key = parse_s3path(s3path) - - s3_ak, s3_sk, s3_endpoint = get_s3_config(bucket) - s3_rw = S3ReaderWriter( - s3_ak, s3_sk, s3_endpoint, "auto", remove_non_official_s3_args(s3path) - ) - may_range_params = parse_s3_range_params(s3path) - if may_range_params is None or 2 != len(may_range_params): - byte_start, byte_end = 0, None - else: - byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1]) - byte_end += byte_start - 1 - return s3_rw.read_jsonl( - remove_non_official_s3_args(s3path), - byte_start, - byte_end, - AbsReaderWriter.MODE_BIN, - ) - - jso = json_parse.loads(read_s3_path(json).decode("utf-8")) - s3_file_path = jso.get("file_location") - if s3_file_path is None: - s3_file_path = jso.get("path") - pdf_file_name = Path(s3_file_path).stem - pdf_data = read_s3_path(s3_file_path) - - do_parse( - pdf_file_name, - pdf_data, - jso["doc_layout_result"], - method, - ) - - -@cli.command() -@click.option("--local_json", type=str, help="输入一个本地jsonl路径") -@click.option( - "--method", - type=parse_pdf_methods, - help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法", - default="auto", -) -@click.option("--inside_model", type=click.BOOL, default=True, help="使用内置模型测试") -@click.option("--model_mode", type=click.STRING, default="full", - help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢") -def local_json_command(local_json, method, inside_model, model_mode): - model_config.__use_inside_model__ = inside_model - model_config.__model_mode__ = model_mode - - def read_s3_path(s3path): - bucket, key = parse_s3path(s3path) - - s3_ak, s3_sk, s3_endpoint = get_s3_config(bucket) - s3_rw = S3ReaderWriter( - s3_ak, s3_sk, s3_endpoint, "auto", remove_non_official_s3_args(s3path) - ) - may_range_params = parse_s3_range_params(s3path) - if may_range_params is None or 2 != len(may_range_params): - byte_start, byte_end = 0, None - else: - byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1]) - byte_end += byte_start - 1 - return s3_rw.read_jsonl( - remove_non_official_s3_args(s3path), - byte_start, - byte_end, - AbsReaderWriter.MODE_BIN, - ) - - with open(local_json, "r", encoding="utf-8") as f: - for json_line in f: - jso = json_parse.loads(json_line) - - s3_file_path = jso.get("file_location") - if s3_file_path is None: - s3_file_path = jso.get("path") - pdf_file_name = Path(s3_file_path).stem - pdf_data = read_s3_path(s3_file_path) - do_parse( - pdf_file_name, - pdf_data, - jso["doc_layout_result"], - method, - ) - - -@cli.command() -@click.option( - "--pdf", type=click.Path(exists=True), required=True, - help='pdf 文件路径, 支持单个文件或文件列表, 文件列表需要以".list"结尾, 一行一个pdf文件路径') -@click.option("--model", type=click.Path(exists=True), help="模型的路径") -@click.option( - "--method", - type=parse_pdf_methods, - help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法", - default="auto", -) -@click.option("--inside_model", type=click.BOOL, default=True, help="使用内置模型测试") -@click.option("--model_mode", type=click.STRING, default="full", - help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢") -def pdf_command(pdf, model, method, inside_model, model_mode): - model_config.__use_inside_model__ = inside_model - model_config.__model_mode__ = model_mode - - def read_fn(path): - disk_rw = DiskReaderWriter(os.path.dirname(path)) - return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN) - - def get_model_json(model_path, doc_path): - # 这里处理pdf和模型相关的逻辑 - if model_path is None: - file_name_without_extension, extension = os.path.splitext(doc_path) - if extension == ".pdf": - model_path = file_name_without_extension + ".json" - else: - raise Exception("pdf_path input error") - if not os.path.exists(model_path): - logger.warning( - f"not found json {model_path} existed" - ) - # 本地无模型数据则调用内置paddle分析，先传空list，在内部识别到空list再调用paddle - model_json = "[]" - else: - model_json = read_fn(model_path).decode("utf-8") - else: - model_json = read_fn(model_path).decode("utf-8") - - return model_json - - def parse_doc(doc_path): - try: - file_name = str(Path(doc_path).stem) - pdf_data = read_fn(doc_path) - jso = json_parse.loads(get_model_json(model, doc_path)) - - do_parse( - file_name, - pdf_data, - jso, - method, - ) - - except Exception as e: - logger.exception(e) - - if not pdf: - logger.error(f"Error: Missing argument '--pdf'.") - exit(f"Error: Missing argument '--pdf'.") - else: - '''适配多个文档的list文件输入''' - if pdf.endswith(".list"): - with open(pdf, "r") as f: - for line in f.readlines(): - line = line.strip() - parse_doc(line) - else: - '''适配单个文档的输入''' - parse_doc(pdf) - - -if __name__ == "__main__": - """ - python magic_pdf/cli/magicpdf.py json-command --json s3://llm-pdf-text/pdf_ebook_and_paper/manual/v001/part-660407a28beb-000002.jsonl?bytes=0,63551 - """ - cli() diff --git a/magic_pdf/dict2md/ocr_mkcontent.py b/magic_pdf/dict2md/ocr_mkcontent.py index 8702d054..2195f7c5 100644 --- a/magic_pdf/dict2md/ocr_mkcontent.py +++ b/magic_pdf/dict2md/ocr_mkcontent.py @@ -168,7 +168,7 @@ def merge_para_with_text(para_block): else: content = ocr_escape_special_markdown_char(content) elif span_type == ContentType.InlineEquation: - content = f"${span['content']}$" + content = f" ${span['content']}$ " elif span_type == ContentType.InterlineEquation: content = f"\n$$\n{span['content']}\n$$\n" diff --git a/magic_pdf/libs/config_reader.py b/magic_pdf/libs/config_reader.py index 34817efc..8006a4ab 100644 --- a/magic_pdf/libs/config_reader.py +++ b/magic_pdf/libs/config_reader.py @@ -57,16 +57,6 @@ def get_bucket_name(path): return bucket -def get_local_dir(): - config = read_config() - local_dir = config.get("temp-output-dir") - if local_dir is None: - logger.warning(f"'temp-output-dir' not found in {CONFIG_FILE_NAME}, use '/tmp' as default") - return "/tmp" - else: - return local_dir - - def get_local_models_dir(): config = read_config() models_dir = config.get("models-dir") @@ -92,5 +82,15 @@ def get_table_recog_config(): return table_config +def get_table_recog_config(): + config = read_config() + table_config = config.get("table-config") + if table_config is None: + logger.warning(f"'table-config' not found in {CONFIG_FILE_NAME}, use 'False' as default") + return json.loads('{"is_table_recog_enable": false, "max_time": 400}') + else: + return table_config + + if __name__ == "__main__": ak, sk, endpoint = get_s3_config("llm-raw") diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index f4e6f9cd..aa5c06a1 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -3,6 +3,7 @@ import os import time from pypandoc import convert_text + os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新 try: import cv2 @@ -107,8 +108,8 @@ class CustomPEKModel: self.apply_table = self.table_config.get("is_table_recog_enable", False) self.apply_ocr = ocr logger.info( - "DocAnalysis init, this may take some times. apply_layout: {}, apply_formula: {}, apply_ocr: {}".format( - self.apply_layout, self.apply_formula, self.apply_ocr + "DocAnalysis init, this may take some times. apply_layout: {}, apply_formula: {}, apply_ocr: {}, apply_table: {}".format( + self.apply_layout, self.apply_formula, self.apply_ocr, self.apply_table ) ) assert self.apply_layout, "DocAnalysis must contain layout model." diff --git a/magic_pdf/pdf_parse_for_train.py b/magic_pdf/pdf_parse_for_train.py deleted file mode 100644 index 58d88df7..00000000 --- a/magic_pdf/pdf_parse_for_train.py +++ /dev/null @@ -1,685 +0,0 @@ -import time - -# from anyio import Path - -from magic_pdf.libs.commons import ( - fitz, - get_delta_time, - get_img_s3_client, - get_docx_model_output, -) -import json -import os -from copy import deepcopy -import math -from loguru import logger -from magic_pdf.layout.bbox_sort import ( - prepare_bboxes_for_layout_split, -) -from magic_pdf.layout.layout_sort import ( - LAYOUT_UNPROC, - get_bboxes_layout, - get_columns_cnt_of_layout, - sort_text_block, -) -from magic_pdf.libs.drop_reason import DropReason -from magic_pdf.libs.markdown_utils import escape_special_markdown_char -from magic_pdf.libs.safe_filename import sanitize_filename -from magic_pdf.libs.vis_utils import draw_bbox_on_page, draw_layout_bbox_on_page -from magic_pdf.pre_proc.cut_image import txt_save_images_by_bboxes -from magic_pdf.pre_proc.detect_images import parse_images -from magic_pdf.pre_proc.detect_tables import parse_tables # 获取tables的bbox -from magic_pdf.pre_proc.detect_equation import parse_equations # 获取equations的bbox -from magic_pdf.pre_proc.detect_header import parse_headers # 获取headers的bbox -from magic_pdf.pre_proc.detect_page_number import parse_pageNos # 获取pageNos的bbox -from magic_pdf.pre_proc.detect_footnote import ( - parse_footnotes_by_model, - parse_footnotes_by_rule, -) # 获取footnotes的bbox -from magic_pdf.pre_proc.detect_footer_by_model import parse_footers # 获取footers的bbox - -from magic_pdf.post_proc.detect_para import ( - ParaProcessPipeline, - TitleDetectionException, - TitleLevelException, - ParaSplitException, - ParaMergeException, - DenseSingleLineBlockException, -) -from magic_pdf.pre_proc.main_text_font import get_main_text_font -from magic_pdf.pre_proc.remove_colored_strip_bbox import remove_colored_strip_textblock -from magic_pdf.pre_proc.remove_footer_header import remove_headder_footer_one_page -from magic_pdf.train_utils.extract_caption import extract_caption_bbox - -""" -from para.para_pipeline import ParaProcessPipeline -from para.exceptions import ( - TitleDetectionException, - TitleLevelException, - ParaSplitException, - ParaMergeException, - DenseSingleLineBlockException, -) -""" - -from magic_pdf.libs.commons import read_file, join_path -from magic_pdf.post_proc.remove_footnote import ( - merge_footnote_blocks, - remove_footnote_blocks, -) -from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker -from magic_pdf.pre_proc.equations_replace import ( - combine_chars_to_pymudict, - remove_chars_in_text_blocks, - replace_equations_in_textblock, -) -from magic_pdf.pre_proc.pdf_pre_filter import pdf_filter -from magic_pdf.pre_proc.detect_footer_header_by_statistics import drop_footer_header -from magic_pdf.pre_proc.construct_page_dict import construct_page_component -from magic_pdf.pre_proc.fix_image import ( - combine_images, - fix_image_vertical, - fix_seperated_image, - include_img_title, -) -from magic_pdf.post_proc.pdf_post_filter import pdf_post_filter -from magic_pdf.pre_proc.remove_rotate_bbox import ( - get_side_boundry, - remove_rotate_side_textblock, - remove_side_blank_block, -) -from magic_pdf.pre_proc.resolve_bbox_conflict import ( - check_text_block_horizontal_overlap, - resolve_bbox_overlap_conflict, -) -from magic_pdf.pre_proc.fix_table import ( - fix_table_text_block, - fix_tables, - include_table_title, -) -from magic_pdf.pre_proc.solve_line_alien import solve_inline_too_large_interval - -denseSingleLineBlockException_msg = DenseSingleLineBlockException().message -titleDetectionException_msg = TitleDetectionException().message -titleLevelException_msg = TitleLevelException().message -paraSplitException_msg = ParaSplitException().message -paraMergeException_msg = ParaMergeException().message - - -def parse_pdf_for_train( - s3_pdf_path, - s3_pdf_profile, - pdf_model_output, - save_path, - book_name, - image_s3_config=None, - start_page_id=0, - end_page_id=None, - junk_img_bojids=[], - debug_mode=False, -): - pdf_bytes = read_file(s3_pdf_path, s3_pdf_profile) - save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest") - md_bookname_save_path = "" - book_name = sanitize_filename(book_name) - if debug_mode: - save_path = join_path(save_tmp_path, "md") - pdf_local_path = join_path(save_tmp_path, "download-pdfs", book_name) - - if not os.path.exists(os.path.dirname(pdf_local_path)): - # 如果目录不存在，创建它 - os.makedirs(os.path.dirname(pdf_local_path)) - - md_bookname_save_path = join_path(save_tmp_path, "md", book_name) - if not os.path.exists(md_bookname_save_path): - # 如果目录不存在，创建它 - os.makedirs(md_bookname_save_path) - - with open(pdf_local_path + ".pdf", "wb") as pdf_file: - pdf_file.write(pdf_bytes) - - pdf_docs = fitz.open("pdf", pdf_bytes) - pdf_info_dict = {} - img_s3_client = get_img_s3_client( - save_path, image_s3_config - ) # 更改函数名和参数，避免歧义 - # img_s3_client = "img_s3_client" #不创建这个对象，直接用字符串占位 - - start_time = time.time() - - """通过统计pdf全篇文字,识别正文字体""" - main_text_font = get_main_text_font(pdf_docs) - - 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): - page = pdf_docs[page_id] - page_width = page.rect.width - page_height = page.rect.height - - if debug_mode: - time_now = time.time() - logger.info( - f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}" - ) - start_time = time_now - """ - # 通过一个规则，过滤掉单页超过1500非junkimg的pdf - # 对单页面非重复id的img数量做统计,如果当前页超过1500则直接return need_drop - """ - page_imgs = page.get_images() - img_counts = 0 - for img in page_imgs: - img_bojid = img[0] - if img_bojid in junk_img_bojids: # 判断这个图片在不在junklist中 - continue # 如果在junklist就不用管了，跳过 - else: - recs = page.get_image_rects(img, transform=True) - if recs: # 如果这张图在当前页面有展示 - img_counts += 1 - if ( - img_counts >= 1500 - ): # 如果去除了junkimg的影响，单页img仍然超过1500的话，就排除当前pdf - logger.warning( - f"page_id: {page_id}, img_counts: {img_counts}, drop this pdf: {book_name}, drop_reason: {DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS}" - ) - result = { - "_need_drop": True, - "_drop_reason": DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS, - } - if not debug_mode: - return result - - """ - ================================================================================================================================== - 首先获取基本的block数据，对pdf进行分解，获取图片、表格、公式、text的bbox - """ - # 解析pdf原始文本block - text_raw_blocks = page.get_text( - "dict", - flags=fitz.TEXTFLAGS_TEXT, - )["blocks"] - model_output_json = get_docx_model_output( - pdf_model_output, page_id - ) - - # 解析图片 - image_bboxes = parse_images(page_id, page, model_output_json, junk_img_bojids) - image_bboxes = fix_image_vertical( - image_bboxes, text_raw_blocks - ) # 修正图片的位置 - image_bboxes = fix_seperated_image(image_bboxes) # 合并有边重合的图片 - - old_image_bboxes = deepcopy(image_bboxes) - image_bboxes = include_img_title( - text_raw_blocks, image_bboxes - ) # 向图片上方和下方寻找title，使用规则进行匹配，暂时只支持英文规则 - """此时image_bboxes中可能出现这种情况，水平并列的2个图片，下方分别有各自的子标题，2个子标题下方又有大标题（形如Figxxx)，会出现2个图片的bbox都包含了这个大标题，这种情况需要把图片合并""" - image_bboxes = combine_images(image_bboxes) # 合并图片 - - # 解析表格并对table_bboxes进行位置的微调,防止表格周围的文字被截断 - table_bboxes = parse_tables(page_id, page, model_output_json) - table_bboxes = fix_tables( - page, table_bboxes, include_table_title=False, scan_line_num=2 - ) # 修正 - table_bboxes = fix_table_text_block( - text_raw_blocks, table_bboxes - ) # 修正与text block的关系,某些table修正与pymupdf获取到的table内textblock没有完全包含，因此要进行一次修正。 - # debug_show_bbox(pdf_docs, page_id, table_bboxes, [], [b['bbox'] for b in text_raw_blocks], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 7) - - old_table_bboxes = deepcopy(table_bboxes) - table_bboxes = include_table_title( - text_raw_blocks, table_bboxes - ) # 向table上方和下方寻找title，使用规则进行匹配，暂时只支持英文规则 - - # 解析公式 - equations_inline_bboxes, equations_interline_bboxes = parse_equations( - page_id, page, model_output_json - ) - - # get image box and caption ! - image_bboxes_with_caption = extract_caption_bbox(image_bboxes, old_image_bboxes) - - # get table box and caption ! - table_bboxes_with_caption = extract_caption_bbox(table_bboxes, old_table_bboxes) - - """ - ================================================================================================================================== - 进入预处理-1阶段 - ------------------- - # # 解析标题 - # title_bboxs = parse_titles(page_id, page, model_output_json) - # # 评估Layout是否规整、简单 - # isSimpleLayout_flag, fullColumn_cnt, subColumn_cnt, curPage_loss = evaluate_pdf_layout(page_id, page, model_output_json) - 接下来开始进行预处理过程 - """ - # title_bboxs = parse_titles(page_id, page, model_output_json) - - """去掉每页的页码、页眉、页脚""" - page_no_bboxs = parse_pageNos(page_id, page, model_output_json) - header_bboxs = parse_headers(page_id, page, model_output_json) - footer_bboxs = parse_footers(page_id, page, model_output_json) - ( - image_bboxes, - table_bboxes, - remain_text_blocks, - removed_hdr_foot_txt_block, - removed_hdr_foot_img_block, - removed_hdr_foot_table, - ) = remove_headder_footer_one_page( - text_raw_blocks, - image_bboxes, - table_bboxes, - header_bboxs, - footer_bboxs, - page_no_bboxs, - page_width, - page_height, - ) - - """去除页面上半部分长条色块内的文本块""" - remain_text_blocks, removed_colored_narrow_strip_background_text_block = ( - remove_colored_strip_textblock(remain_text_blocks, page) - ) - - # debug_show_bbox(pdf_docs, page_id, footnote_bboxes_by_model, [b['bbox'] for b in remain_text_blocks], header_bboxs, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 7) - - """去掉旋转的文字：水印、垂直排列的文字""" - remain_text_blocks, removed_non_horz_text_block = remove_rotate_side_textblock( - remain_text_blocks, page_width, page_height - ) # 去掉水印，非水平文字 - remain_text_blocks, removed_empty_side_block = remove_side_blank_block( - remain_text_blocks, page_width, page_height - ) # 删除页面四周可能会留下的完全空白的textblock，这种block形成原因未知 - - """出现在图片、表格上的文字块去掉，把层叠的图片单独分离出来，不参与layout的计算""" - ( - image_bboxes, - table_bboxes, - equations_interline_bboxes, - equations_inline_bboxes, - remain_text_blocks, - text_block_on_image_removed, - images_overlap_backup, - interline_eq_temp_text_block, - ) = resolve_bbox_overlap_conflict( - image_bboxes, - table_bboxes, - equations_interline_bboxes, - equations_inline_bboxes, - remain_text_blocks, - ) - - # """去掉footnote, 从文字和图片中""" - # # 通过模型识别到的footnote - # 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) - """ - ================================================================================================================================== - """ - if debug_mode: # debugmode截图到本地 - save_path = join_path(save_tmp_path, "md") - - # 把图、表、公式都进行截图，保存到存储上，返回图片路径作为内容 - image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info = ( - txt_save_images_by_bboxes( - book_name, - page_id, - page, - save_path, - image_bboxes, - images_overlap_backup, - table_bboxes, - equations_inline_bboxes, - equations_interline_bboxes, - # 传入img_s3_client - img_s3_client, - ) - ) # 只要表格和图片的截图 - - """"以下进入到公式替换环节 """ - char_level_text_blocks = page.get_text("rawdict", flags=fitz.TEXTFLAGS_TEXT)[ - "blocks" - ] - remain_text_blocks = combine_chars_to_pymudict( - remain_text_blocks, char_level_text_blocks - ) # 合并chars - remain_text_blocks = replace_equations_in_textblock( - remain_text_blocks, inline_eq_info, interline_eq_info - ) - remain_text_blocks = remove_citation_marker( - remain_text_blocks - ) # 公式替换之后去角标，防止公式无法替换成功。但是这样也会带来个问题就是把角标当公式。各有优劣。 - remain_text_blocks = remove_chars_in_text_blocks( - remain_text_blocks - ) # 减少中间态数据体积 - # debug_show_bbox(pdf_docs, page_id, [b['bbox'] for b in inline_eq_info], [b['bbox'] for b in interline_eq_info], [], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 3) - - """去掉footnote, 从文字和图片中(先去角标再去footnote试试)""" - # 通过模型识别到的footnote - 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: - logger.warning( - f"page_id: {page_id}, drop this pdf: {book_name}, reason: {err}" - ) - if not debug_mode: - return err - - """ - ================================================================================================================================== - 进行版面布局切分和过滤 - """ - """在切分之前，先检查一下bbox是否有左右重叠的情况，如果有，那么就认为这个pdf暂时没有能力处理好，这种左右重叠的情况大概率是由于pdf里的行间公式、表格没有被正确识别出来造成的 """ - - is_text_block_horz_overlap = check_text_block_horizontal_overlap( - remain_text_blocks, header_bboxs, footer_bboxs - ) - - if is_text_block_horz_overlap: - # debug_show_bbox(pdf_docs, page_id, [b['bbox'] for b in remain_text_blocks], [], [], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 0) - logger.warning( - f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TEXT_BLCOK_HOR_OVERLAP}" - ) - result = { - "_need_drop": True, - "_drop_reason": DropReason.TEXT_BLCOK_HOR_OVERLAP, - } - if not debug_mode: - return result - - """统一格式化成一个数据结构用于计算layout""" - page_y0 = 0 if len(header_bboxs) == 0 else max([b[3] for b in header_bboxs]) - page_y1 = ( - page_height if len(footer_bboxs) == 0 else min([b[1] for b in footer_bboxs]) - ) - left_x, right_x = get_side_boundry( - removed_non_horz_text_block, page_width, page_height - ) - page_boundry = [ - math.floor(left_x), - page_y0 + 1, - math.ceil(right_x), - page_y1 - 1, - ] - # 返回的是一个数组，每个元素[x0, y0, x1, y1, block_content, idx_x, idx_y], 初始时候idx_x, idx_y都是None. 对于图片、公式来说，block_content是图片的地址， 对于段落来说，block_content是段落的内容 - - all_bboxes = prepare_bboxes_for_layout_split( - image_info, - image_backup_info, - table_info, - inline_eq_info, - interline_eq_info, - remain_text_blocks, - page_boundry, - page, - ) - # debug_show_bbox(pdf_docs, page_id, [], [], all_bboxes, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 1) - """page_y0, page_y1能够过滤掉页眉和页脚，不会算作layout内""" - layout_bboxes, layout_tree = get_bboxes_layout( - all_bboxes, page_boundry, page_id - ) - - if ( - len(remain_text_blocks) > 0 - and len(all_bboxes) > 0 - and len(layout_bboxes) == 0 - ): - logger.warning( - f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.CAN_NOT_DETECT_PAGE_LAYOUT}" - ) - result = { - "_need_drop": True, - "_drop_reason": DropReason.CAN_NOT_DETECT_PAGE_LAYOUT, - } - if not debug_mode: - return result - - """以下去掉复杂的布局和超过2列的布局""" - if any( - [lay["layout_label"] == LAYOUT_UNPROC for lay in layout_bboxes] - ): # 复杂的布局 - logger.warning( - f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.COMPLICATED_LAYOUT}" - ) - result = {"_need_drop": True, "_drop_reason": DropReason.COMPLICATED_LAYOUT} - if not debug_mode: - return result - - layout_column_width = get_columns_cnt_of_layout(layout_tree) - if layout_column_width > 2: # 去掉超过2列的布局pdf - logger.warning( - f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TOO_MANY_LAYOUT_COLUMNS}" - ) - result = { - "_need_drop": True, - "_drop_reason": DropReason.TOO_MANY_LAYOUT_COLUMNS, - "extra_info": {"column_cnt": layout_column_width}, - } - if not debug_mode: - return result - - """ - ================================================================================================================================== - 构造出下游需要的数据结构 - """ - remain_text_blocks = ( - remain_text_blocks + interline_eq_temp_text_block - ) # 把计算layout时候临时删除的行间公式再放回去，防止行间公式替换的时候丢失。 - removed_text_blocks = [] - removed_text_blocks.extend(removed_hdr_foot_txt_block) - # removed_text_blocks.extend(removed_footnote_text_block) - removed_text_blocks.extend(text_block_on_image_removed) - removed_text_blocks.extend(removed_non_horz_text_block) - removed_text_blocks.extend(removed_colored_narrow_strip_background_text_block) - - removed_images = [] - # removed_images.extend(footnote_imgs) - removed_images.extend(removed_hdr_foot_img_block) - - images_backup = [] - images_backup.extend(image_backup_info) - remain_text_blocks = escape_special_markdown_char( - remain_text_blocks - ) # 转义span里的text - sorted_text_remain_text_block = sort_text_block( - remain_text_blocks, layout_bboxes - ) - - footnote_bboxes_tmp = [] - footnote_bboxes_tmp.extend(footnote_bboxes_by_model) - footnote_bboxes_tmp.extend(footnote_bboxes_by_rule) - - page_info = construct_page_component( - page_id, - image_info, - table_info, - sorted_text_remain_text_block, - layout_bboxes, - inline_eq_info, - interline_eq_info, - page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"], - removed_text_blocks=removed_text_blocks, - removed_image_blocks=removed_images, - images_backup=images_backup, - droped_table_block=[], - table_backup=[], - layout_tree=layout_tree, - page_w=page.rect.width, - page_h=page.rect.height, - footnote_bboxes_tmp=footnote_bboxes_tmp, - ) - - page_info["image_bboxes_with_caption"] = image_bboxes_with_caption # add by xr - page_info["table_bboxes_with_caption"] = table_bboxes_with_caption - - page_info["bak_page_no_bboxes"] = page_no_bboxs - page_info["bak_header_bboxes"] = header_bboxs - page_info["bak_footer_bboxes"] = footer_bboxs - page_info["bak_footer_note_bboxes"] = footnote_bboxes_tmp - - pdf_info_dict[f"page_{page_id}"] = page_info - - # end page for - - """计算后处理阶段耗时""" - start_time = time.time() - - """ - ================================================================================================================================== - 去掉页眉和页脚，这里需要用到一定的统计量，所以放到最后 - 页眉和页脚主要从文本box和图片box中去除，位于页面的四周。 - 下面函数会直接修改pdf_info_dict,从文字块中、图片中删除属于页眉页脚的内容，删除内容做相对应记录 - """ - # 去页眉页脚 - header, footer = drop_footer_header( - pdf_info_dict - ) # TODO: using header and footer boxes here ! - - """对单个layout内footnote和他下面的所有textbbox合并""" - - for page_key, page_info in pdf_info_dict.items(): - page_info = merge_footnote_blocks(page_info, main_text_font) - page_info = remove_footnote_blocks(page_info) - pdf_info_dict[page_key] = page_info - - """进入pdf后置过滤器，去掉一些不合理的pdf""" - - i = 0 - for page_info in pdf_info_dict.values(): - is_good_pdf, err = pdf_post_filter(page_info) - if not is_good_pdf: - logger.warning(f"page_id: {i}, drop this pdf: {book_name}, reason: {err}") - if not debug_mode: - return err - i += 1 - - if debug_mode: - params_file_save_path = join_path( - save_tmp_path, "md", book_name, "preproc_out.json" - ) - page_draw_rect_save_path = join_path( - save_tmp_path, "md", book_name, "layout.pdf" - ) - # dir_path = os.path.dirname(page_draw_rect_save_path) - # if not os.path.exists(dir_path): - # # 如果目录不存在，创建它 - # os.makedirs(dir_path) - - with open(params_file_save_path, "w", encoding="utf-8") as f: - json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4) - # 先检测本地 page_draw_rect_save_path 是否存在，如果存在则删除 - if os.path.exists(page_draw_rect_save_path): - os.remove(page_draw_rect_save_path) - # 绘制bbox和layout到pdf - draw_bbox_on_page(pdf_docs, pdf_info_dict, page_draw_rect_save_path) - draw_layout_bbox_on_page( - pdf_docs, pdf_info_dict, header, footer, page_draw_rect_save_path - ) - - if debug_mode: - # 打印后处理阶段耗时 - logger.info(f"post_processing_time: {get_delta_time(start_time)}") - - """ - ================================================================================================================================== - 进入段落处理-2阶段 - """ - - # 处理行内文字间距较大问题 - pdf_info_dict = solve_inline_too_large_interval(pdf_info_dict) - - start_time = time.time() - - para_process_pipeline = ParaProcessPipeline() - - def _deal_with_text_exception(error_info): - logger.warning( - f"page_id: {page_id}, drop this pdf: {book_name}, reason: {error_info}" - ) - if error_info == denseSingleLineBlockException_msg: - logger.warning( - f"Drop this pdf: {book_name}, reason: {DropReason.DENSE_SINGLE_LINE_BLOCK}" - ) - result = { - "_need_drop": True, - "_drop_reason": DropReason.DENSE_SINGLE_LINE_BLOCK, - } - return result - if error_info == titleDetectionException_msg: - logger.warning( - f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_DETECTION_FAILED}" - ) - result = { - "_need_drop": True, - "_drop_reason": DropReason.TITLE_DETECTION_FAILED, - } - return result - elif error_info == titleLevelException_msg: - logger.warning( - f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_LEVEL_FAILED}" - ) - result = {"_need_drop": True, "_drop_reason": DropReason.TITLE_LEVEL_FAILED} - return result - elif error_info == paraSplitException_msg: - logger.warning( - f"Drop this pdf: {book_name}, reason: {DropReason.PARA_SPLIT_FAILED}" - ) - result = {"_need_drop": True, "_drop_reason": DropReason.PARA_SPLIT_FAILED} - return result - elif error_info == paraMergeException_msg: - logger.warning( - f"Drop this pdf: {book_name}, reason: {DropReason.PARA_MERGE_FAILED}" - ) - result = {"_need_drop": True, "_drop_reason": DropReason.PARA_MERGE_FAILED} - return result - - if debug_mode: - input_pdf_file = f"{pdf_local_path}.pdf" - output_dir = f"{save_path}/{book_name}" - output_pdf_file = f"{output_dir}/pdf_annos.pdf" - - """ - Call the para_process_pipeline function to process the pdf_info_dict. - - Parameters: - para_debug_mode: str or None - If para_debug_mode is None, the para_process_pipeline will not keep any intermediate results. - If para_debug_mode is "simple", the para_process_pipeline will only keep the annos on the pdf and the final results as a json file. - If para_debug_mode is "full", the para_process_pipeline will keep all the intermediate results generated during each step. - """ - pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline( - pdf_info_dict, - para_debug_mode="simple", - input_pdf_path=input_pdf_file, - output_pdf_path=output_pdf_file, - ) - # 打印段落处理阶段耗时 - logger.info(f"para_process_time: {get_delta_time(start_time)}") - - # debug的时候不return drop信息 - if error_info is not None: - _deal_with_text_exception(error_info) - return pdf_info_dict - else: - pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline( - pdf_info_dict - ) - if error_info is not None: - return _deal_with_text_exception(error_info) - - return pdf_info_dict diff --git a/magic_pdf/rw/AbsReaderWriter.py b/magic_pdf/rw/AbsReaderWriter.py index ce5e0774..78e455af 100644 --- a/magic_pdf/rw/AbsReaderWriter.py +++ b/magic_pdf/rw/AbsReaderWriter.py @@ -2,33 +2,16 @@ from abc import ABC, abstractmethod class AbsReaderWriter(ABC): - """ - 同时支持二进制和文本读写的抽象类 - """ MODE_TXT = "text" MODE_BIN = "binary" - - def __init__(self, parent_path): - # 初始化代码可以在这里添加，如果需要的话 - self.parent_path = parent_path # 对于本地目录是父目录，对于s3是会写到这个path下。 - @abstractmethod def read(self, path: str, mode=MODE_TXT): - """ - 无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path - """ raise NotImplementedError @abstractmethod def write(self, content: str, path: str, mode=MODE_TXT): - """ - 无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path - """ raise NotImplementedError @abstractmethod - def read_jsonl(self, path: str, byte_start=0, byte_end=None, encoding='utf-8'): - """ - 无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path - """ + def read_offset(self, path: str, offset=0, limit=None) -> bytes: raise NotImplementedError diff --git a/magic_pdf/rw/DiskReaderWriter.py b/magic_pdf/rw/DiskReaderWriter.py index 4e034c80..dd627fde 100644 --- a/magic_pdf/rw/DiskReaderWriter.py +++ b/magic_pdf/rw/DiskReaderWriter.py @@ -3,34 +3,29 @@ from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter from loguru import logger -MODE_TXT = "text" -MODE_BIN = "binary" - - class DiskReaderWriter(AbsReaderWriter): - def __init__(self, parent_path, encoding="utf-8"): self.path = parent_path self.encoding = encoding - def read(self, path, mode=MODE_TXT): + def read(self, path, mode=AbsReaderWriter.MODE_TXT): if os.path.isabs(path): abspath = path else: abspath = os.path.join(self.path, path) if not os.path.exists(abspath): - logger.error(f"文件 {abspath} 不存在") - raise Exception(f"文件 {abspath} 不存在") - if mode == MODE_TXT: + logger.error(f"file {abspath} not exists") + raise Exception(f"file {abspath} no exists") + if mode == AbsReaderWriter.MODE_TXT: with open(abspath, "r", encoding=self.encoding) as f: return f.read() - elif mode == MODE_BIN: + elif mode == AbsReaderWriter.MODE_BIN: with open(abspath, "rb") as f: return f.read() else: raise ValueError("Invalid mode. Use 'text' or 'binary'.") - def write(self, content, path, mode=MODE_TXT): + def write(self, content, path, mode=AbsReaderWriter.MODE_TXT): if os.path.isabs(path): abspath = path else: @@ -38,29 +33,42 @@ class DiskReaderWriter(AbsReaderWriter): directory_path = os.path.dirname(abspath) if not os.path.exists(directory_path): os.makedirs(directory_path) - if mode == MODE_TXT: + if mode == AbsReaderWriter.MODE_TXT: with open(abspath, "w", encoding=self.encoding, errors="replace") as f: f.write(content) - elif mode == MODE_BIN: + elif mode == AbsReaderWriter.MODE_BIN: with open(abspath, "wb") as f: f.write(content) else: raise ValueError("Invalid mode. Use 'text' or 'binary'.") - def read_jsonl(self, path: str, byte_start=0, byte_end=None, encoding="utf-8"): - return self.read(path) + def read_offset(self, path: str, offset=0, limit=None): + abspath = path + if not os.path.isabs(path): + abspath = os.path.join(self.path, path) + with open(abspath, "rb") as f: + f.seek(offset) + return f.read(limit) -# 使用示例 if __name__ == "__main__": - file_path = "io/test/example.txt" - drw = DiskReaderWriter("D:\projects\papayfork\Magic-PDF\magic_pdf") + if 0: + file_path = "io/test/example.txt" + drw = DiskReaderWriter("D:\projects\papayfork\Magic-PDF\magic_pdf") - # 写入内容到文件 - drw.write(b"Hello, World!", path="io/test/example.txt", mode="binary") + # 写入内容到文件 + drw.write(b"Hello, World!", path="io/test/example.txt", mode="binary") + + # 从文件读取内容 + content = drw.read(path=file_path) + if content: + logger.info(f"从 {file_path} 读取的内容: {content}") + if 1: + drw = DiskReaderWriter("/opt/data/pdf/resources/test/io/") + content_bin = drw.read_offset("1.txt") + assert content_bin == b"ABCD!" + + content_bin = drw.read_offset("1.txt", offset=1, limit=2) + assert content_bin == b"BC" - # 从文件读取内容 - content = drw.read(path=file_path) - if content: - logger.info(f"从 {file_path} 读取的内容: {content}") diff --git a/magic_pdf/rw/S3ReaderWriter.py b/magic_pdf/rw/S3ReaderWriter.py index df394a1e..ac1cc1f1 100644 --- a/magic_pdf/rw/S3ReaderWriter.py +++ b/magic_pdf/rw/S3ReaderWriter.py @@ -2,16 +2,18 @@ from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter from magic_pdf.libs.commons import parse_aws_param, parse_bucket_key, join_path import boto3 from loguru import logger -from boto3.s3.transfer import TransferConfig from botocore.config import Config -import os - -MODE_TXT = "text" -MODE_BIN = "binary" class S3ReaderWriter(AbsReaderWriter): - def __init__(self, ak: str, sk: str, endpoint_url: str, addressing_style: str = 'auto', parent_path: str = ''): + def __init__( + self, + ak: str, + sk: str, + endpoint_url: str, + addressing_style: str = "auto", + parent_path: str = "", + ): self.client = self._get_client(ak, sk, endpoint_url, addressing_style) self.path = parent_path @@ -21,12 +23,14 @@ class S3ReaderWriter(AbsReaderWriter): aws_access_key_id=ak, aws_secret_access_key=sk, endpoint_url=endpoint_url, - config=Config(s3={"addressing_style": addressing_style}, - retries={'max_attempts': 5, 'mode': 'standard'}), + config=Config( + s3={"addressing_style": addressing_style}, + retries={"max_attempts": 5, "mode": "standard"}, + ), ) return s3_client - def read(self, s3_relative_path, mode=MODE_TXT, encoding="utf-8"): + def read(self, s3_relative_path, mode=AbsReaderWriter.MODE_TXT, encoding="utf-8"): if s3_relative_path.startswith("s3://"): s3_path = s3_relative_path else: @@ -34,22 +38,22 @@ class S3ReaderWriter(AbsReaderWriter): bucket_name, key = parse_bucket_key(s3_path) res = self.client.get_object(Bucket=bucket_name, Key=key) body = res["Body"].read() - if mode == MODE_TXT: + if mode == AbsReaderWriter.MODE_TXT: data = body.decode(encoding) # Decode bytes to text - elif mode == MODE_BIN: + elif mode == AbsReaderWriter.MODE_BIN: data = body else: raise ValueError("Invalid mode. Use 'text' or 'binary'.") return data - def write(self, content, s3_relative_path, mode=MODE_TXT, encoding="utf-8"): + def write(self, content, s3_relative_path, mode=AbsReaderWriter.MODE_TXT, encoding="utf-8"): if s3_relative_path.startswith("s3://"): s3_path = s3_relative_path else: s3_path = join_path(self.path, s3_relative_path) - if mode == MODE_TXT: + if mode == AbsReaderWriter.MODE_TXT: body = content.encode(encoding) # Encode text data as bytes - elif mode == MODE_BIN: + elif mode == AbsReaderWriter.MODE_BIN: body = content else: raise ValueError("Invalid mode. Use 'text' or 'binary'.") @@ -57,51 +61,82 @@ class S3ReaderWriter(AbsReaderWriter): self.client.put_object(Body=body, Bucket=bucket_name, Key=key) logger.info(f"内容已写入 {s3_path} ") - def read_jsonl(self, path: str, byte_start=0, byte_end=None, mode=MODE_TXT, encoding='utf-8'): + def read_offset(self, path: str, offset=0, limit=None) -> bytes: if path.startswith("s3://"): s3_path = path else: s3_path = join_path(self.path, path) bucket_name, key = parse_bucket_key(s3_path) - range_header = f'bytes={byte_start}-{byte_end}' if byte_end else f'bytes={byte_start}-' + range_header = ( + f"bytes={offset}-{offset+limit-1}" if limit else f"bytes={offset}-" + ) res = self.client.get_object(Bucket=bucket_name, Key=key, Range=range_header) - body = res["Body"].read() - if mode == MODE_TXT: - data = body.decode(encoding) # Decode bytes to text - elif mode == MODE_BIN: - data = body - else: - raise ValueError("Invalid mode. Use 'text' or 'binary'.") - return data + return res["Body"].read() if __name__ == "__main__": - # Config the connection info - ak = "" - sk = "" - endpoint_url = "" - addressing_style = "auto" - bucket_name = "" - # Create an S3ReaderWriter object - s3_reader_writer = S3ReaderWriter(ak, sk, endpoint_url, addressing_style, "s3://bucket_name/") + if 0: + # Config the connection info + ak = "" + sk = "" + endpoint_url = "" + addressing_style = "auto" + bucket_name = "" + # Create an S3ReaderWriter object + s3_reader_writer = S3ReaderWriter( + ak, sk, endpoint_url, addressing_style, "s3://bucket_name/" + ) - # Write text data to S3 - text_data = "This is some text data" - s3_reader_writer.write(data=text_data, s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_TXT) + # Write text data to S3 + text_data = "This is some text data" + s3_reader_writer.write( + text_data, + s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", + mode=AbsReaderWriter.MODE_TXT, + ) - # Read text data from S3 - text_data_read = s3_reader_writer.read(s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_TXT) - logger.info(f"Read text data from S3: {text_data_read}") - # Write binary data to S3 - binary_data = b"This is some binary data" - s3_reader_writer.write(data=text_data, s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_BIN) + # Read text data from S3 + text_data_read = s3_reader_writer.read( + s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=AbsReaderWriter.MODE_TXT + ) + logger.info(f"Read text data from S3: {text_data_read}") + # Write binary data to S3 + binary_data = b"This is some binary data" + s3_reader_writer.write( + text_data, + s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", + mode=AbsReaderWriter.MODE_BIN, + ) - # Read binary data from S3 - binary_data_read = s3_reader_writer.read(s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_BIN) - logger.info(f"Read binary data from S3: {binary_data_read}") + # Read binary data from S3 + binary_data_read = s3_reader_writer.read( + s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=AbsReaderWriter.MODE_BIN + ) + logger.info(f"Read binary data from S3: {binary_data_read}") - # Range Read text data from S3 - binary_data_read = s3_reader_writer.read_jsonl(path=f"s3://{bucket_name}/ebook/test/test.json", - byte_start=0, byte_end=10, mode=MODE_BIN) - logger.info(f"Read binary data from S3: {binary_data_read}") + # Range Read text data from S3 + binary_data_read = s3_reader_writer.read_offset( + path=f"s3://{bucket_name}/ebook/test/test.json", offset=0, limit=10 + ) + logger.info(f"Read binary data from S3: {binary_data_read}") + if 1: + import os + import json + + ak = os.getenv("AK", "") + sk = os.getenv("SK", "") + endpoint_url = os.getenv("ENDPOINT", "") + bucket = os.getenv("S3_BUCKET", "") + prefix = os.getenv("S3_PREFIX", "") + key_basename = os.getenv("S3_KEY_BASENAME", "") + s3_reader_writer = S3ReaderWriter( + ak, sk, endpoint_url, "auto", f"s3://{bucket}/{prefix}" + ) + content_bin = s3_reader_writer.read_offset(key_basename) + assert content_bin[:10] == b'{"track_id' + assert content_bin[-10:] == b'r":null}}\n' + + content_bin = s3_reader_writer.read_offset(key_basename, offset=424, limit=426) + jso = json.dumps(content_bin.decode("utf-8")) + print(jso) diff --git a/magic_pdf/cli/__init__.py b/magic_pdf/tools/__init__.py similarity index 100% rename from magic_pdf/cli/__init__.py rename to magic_pdf/tools/__init__.py diff --git a/magic_pdf/tools/cli.py b/magic_pdf/tools/cli.py new file mode 100644 index 00000000..76fe3fa6 --- /dev/null +++ b/magic_pdf/tools/cli.py @@ -0,0 +1,79 @@ +import os +import click +from loguru import logger +from pathlib import Path + +from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter +from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter +import magic_pdf.model as model_config +from magic_pdf.tools.common import parse_pdf_methods, do_parse +from magic_pdf.libs.version import __version__ + + +@click.command() +@click.version_option(__version__, "--version", "-v", help="display the version and exit") +@click.option( + "-p", + "--path", + "path", + type=click.Path(exists=True), + required=True, + help="local pdf filepath or directory", +) +@click.option( + "-o", + "--output-dir", + "output_dir", + type=str, + help="output local directory", + default="", +) +@click.option( + "-m", + "--method", + "method", + type=parse_pdf_methods, + help="""the method for parsing pdf. +ocr: using ocr technique to extract information from pdf. +txt: suitable for the text-based pdf only and outperform ocr. +auto: automatically choose the best method for parsing pdf from ocr and txt. +without method specified, auto will be used by default.""", + default="auto", +) +def cli(path, output_dir, method): + model_config.__use_inside_model__ = True + model_config.__model_mode__ = "full" + if output_dir == "": + if os.path.isdir(path): + output_dir = os.path.join(path, "output") + else: + output_dir = os.path.join(os.path.dirname(path), "output") + + def read_fn(path): + disk_rw = DiskReaderWriter(os.path.dirname(path)) + return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN) + + def parse_doc(doc_path: str): + try: + file_name = str(Path(doc_path).stem) + pdf_data = read_fn(doc_path) + do_parse( + output_dir, + file_name, + pdf_data, + [], + method, + ) + + except Exception as e: + logger.exception(e) + + if os.path.isdir(path): + for doc_path in Path(path).glob("*.pdf"): + parse_doc(doc_path) + else: + parse_doc(path) + + +if __name__ == "__main__": + cli() diff --git a/magic_pdf/tools/cli_dev.py b/magic_pdf/tools/cli_dev.py new file mode 100644 index 00000000..9a8830a5 --- /dev/null +++ b/magic_pdf/tools/cli_dev.py @@ -0,0 +1,156 @@ +import os +import json as json_parse +import click +from pathlib import Path +from magic_pdf.libs.path_utils import ( + parse_s3path, + parse_s3_range_params, + remove_non_official_s3_args, +) +from magic_pdf.libs.config_reader import ( + get_s3_config, +) +from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter +from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter +from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter +import magic_pdf.model as model_config +from magic_pdf.tools.common import parse_pdf_methods, do_parse +from magic_pdf.libs.version import __version__ + + +def read_s3_path(s3path): + bucket, key = parse_s3path(s3path) + + s3_ak, s3_sk, s3_endpoint = get_s3_config(bucket) + s3_rw = S3ReaderWriter( + s3_ak, s3_sk, s3_endpoint, "auto", remove_non_official_s3_args(s3path) + ) + may_range_params = parse_s3_range_params(s3path) + if may_range_params is None or 2 != len(may_range_params): + byte_start, byte_end = 0, None + else: + byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1]) + byte_end += byte_start - 1 + return s3_rw.read_jsonl( + remove_non_official_s3_args(s3path), + byte_start, + byte_end, + AbsReaderWriter.MODE_BIN, + ) + + +@click.group() +@click.version_option(__version__, "--version", "-v", help="显示版本信息") +def cli(): + pass + + +@cli.command() +@click.option( + "-j", + "--jsonl", + "jsonl", + type=str, + help="输入 jsonl 路径，本地或者 s3 上的文件", + required=True, +) +@click.option( + "-m", + "--method", + "method", + type=parse_pdf_methods, + help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法", + default="auto", +) +@click.option( + "-o", + "--output-dir", + "output_dir", + type=str, + help="输出到本地目录", + default="", +) +def jsonl(jsonl, method, output_dir): + print("haha") + model_config.__use_inside_model__ = False + full_jsonl_path = os.path.realpath(jsonl) + if output_dir == "": + output_dir = os.path.join(os.path.dirname(full_jsonl_path), "output") + + if jsonl.startswith("s3://"): + jso = json_parse.loads(read_s3_path(jsonl).decode("utf-8")) + else: + with open(jsonl) as f: + jso = json_parse.loads(f.readline()) + s3_file_path = jso.get("file_location") + if s3_file_path is None: + s3_file_path = jso.get("path") + pdf_file_name = Path(s3_file_path).stem + pdf_data = read_s3_path(s3_file_path) + + + print(pdf_file_name, jso, method) + do_parse( + output_dir, + pdf_file_name, + pdf_data, + jso["doc_layout_result"], + method, + f_dump_content_list=True, + ) + + +@cli.command() +@click.option( + "-p", + "--pdf", + "pdf", + type=click.Path(exists=True), + required=True, + help="本地 PDF 文件", +) +@click.option( + "-j", + "--json", + "json_data", + type=click.Path(exists=True), + required=True, + help="本地模型推理出的 json 数据", +) +@click.option( + "-o", "--output-dir", "output_dir", type=str, help="本地输出目录", default="" +) +@click.option( + "-m", + "--method", + "method", + type=parse_pdf_methods, + help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法", + default="auto", +) +def pdf(pdf, json_data, output_dir, method): + model_config.__use_inside_model__ = False + full_pdf_path = os.path.realpath(pdf) + if output_dir == "": + output_dir = os.path.join(os.path.dirname(full_pdf_path), "output") + + def read_fn(path): + disk_rw = DiskReaderWriter(os.path.dirname(path)) + return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN) + + model_json_list = json_parse.loads(read_fn(json_data).decode("utf-8")) + + file_name = str(Path(full_pdf_path).stem) + pdf_data = read_fn(full_pdf_path) + do_parse( + output_dir, + file_name, + pdf_data, + model_json_list, + method, + f_dump_content_list=True, + ) + + +if __name__ == "__main__": + cli() diff --git a/magic_pdf/tools/common.py b/magic_pdf/tools/common.py new file mode 100644 index 00000000..ac27ef21 --- /dev/null +++ b/magic_pdf/tools/common.py @@ -0,0 +1,119 @@ +import os +import json as json_parse +import copy +import click +from loguru import logger +from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode +from magic_pdf.libs.draw_bbox import draw_layout_bbox, draw_span_bbox +from magic_pdf.pipe.UNIPipe import UNIPipe +from magic_pdf.pipe.OCRPipe import OCRPipe +from magic_pdf.pipe.TXTPipe import TXTPipe +from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter +from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter +import magic_pdf.model as model_config + + +def prepare_env(output_dir, pdf_file_name, method): + local_parent_dir = os.path.join(output_dir, pdf_file_name, method) + + local_image_dir = os.path.join(str(local_parent_dir), "images") + local_md_dir = local_parent_dir + os.makedirs(local_image_dir, exist_ok=True) + os.makedirs(local_md_dir, exist_ok=True) + return local_image_dir, local_md_dir + + +def do_parse( + output_dir, + pdf_file_name, + pdf_bytes, + model_list, + parse_method, + f_draw_span_bbox=True, + f_draw_layout_bbox=True, + f_dump_md=True, + f_dump_middle_json=True, + f_dump_model_json=True, + f_dump_orig_pdf=True, + f_dump_content_list=False, + f_make_md_mode=MakeMode.MM_MD, +): + orig_model_list = copy.deepcopy(model_list) + local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method) + + image_writer, md_writer = DiskReaderWriter(local_image_dir), DiskReaderWriter( + local_md_dir + ) + image_dir = str(os.path.basename(local_image_dir)) + + if parse_method == "auto": + jso_useful_key = {"_pdf_type": "", "model_list": model_list} + pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True) + elif parse_method == "txt": + pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True) + elif parse_method == "ocr": + pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True) + else: + logger.error("unknown parse method") + exit(1) + + pipe.pipe_classify() + + if len(model_list) == 0: + if model_config.__use_inside_model__: + pipe.pipe_analyze() + orig_model_list = copy.deepcopy(pipe.model_list) + else: + logger.error("need model list input") + exit(2) + + pipe.pipe_parse() + pdf_info = pipe.pdf_mid_data["pdf_info"] + if f_draw_layout_bbox: + draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir) + if f_draw_span_bbox: + draw_span_bbox(pdf_info, pdf_bytes, local_md_dir) + + md_content = pipe.pipe_mk_markdown( + image_dir, drop_mode=DropMode.NONE, md_make_mode=f_make_md_mode + ) + if f_dump_md: + md_writer.write( + content=md_content, + path=f"{pdf_file_name}.md", + mode=AbsReaderWriter.MODE_TXT, + ) + + if f_dump_middle_json: + md_writer.write( + content=json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4), + path="middle.json", + mode=AbsReaderWriter.MODE_TXT, + ) + + if f_dump_model_json: + md_writer.write( + content=json_parse.dumps(orig_model_list, ensure_ascii=False, indent=4), + path="model.json", + mode=AbsReaderWriter.MODE_TXT, + ) + + if f_dump_orig_pdf: + md_writer.write( + content=pdf_bytes, + path="origin.pdf", + mode=AbsReaderWriter.MODE_BIN, + ) + + content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE) + if f_dump_content_list: + md_writer.write( + content=json_parse.dumps(content_list, ensure_ascii=False, indent=4), + path="content_list.json", + mode=AbsReaderWriter.MODE_TXT, + ) + + logger.info(f"local output dir is {local_md_dir}") + + +parse_pdf_methods = click.Choice(["ocr", "txt", "auto"]) diff --git a/magic_pdf/train_utils/convert_to_train_format.py b/magic_pdf/train_utils/convert_to_train_format.py deleted file mode 100644 index c50e6b3b..00000000 --- a/magic_pdf/train_utils/convert_to_train_format.py +++ /dev/null @@ -1,65 +0,0 @@ -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 diff --git a/magic_pdf/train_utils/extract_caption.py b/magic_pdf/train_utils/extract_caption.py deleted file mode 100644 index 74e0c51d..00000000 --- a/magic_pdf/train_utils/extract_caption.py +++ /dev/null @@ -1,59 +0,0 @@ -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 diff --git a/magic_pdf/train_utils/remove_footer_header.py b/magic_pdf/train_utils/remove_footer_header.py deleted file mode 100644 index 57bd3e9e..00000000 --- a/magic_pdf/train_utils/remove_footer_header.py +++ /dev/null @@ -1,159 +0,0 @@ -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, - ) diff --git a/magic_pdf/train_utils/vis_utils.py b/magic_pdf/train_utils/vis_utils.py deleted file mode 100644 index 996ae514..00000000 --- a/magic_pdf/train_utils/vis_utils.py +++ /dev/null @@ -1,327 +0,0 @@ -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() diff --git a/setup.py b/setup.py index 1081b908..43c611da 100644 --- a/setup.py +++ b/setup.py @@ -36,12 +36,15 @@ if __name__ == '__main__': "paddlepaddle==3.0.0b1;platform_system=='Linux'", "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", ], - "full": ["unimernet==0.1.6", - "matplotlib", - "ultralytics", - "paddleocr==2.7.3", - "paddlepaddle==3.0.0b1;platform_system=='Linux'", - "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", + "full": ["unimernet==0.1.6", # 0.1.6版本大幅裁剪依赖包范围，推荐使用此版本 + "matplotlib<=3.9.0;platform_system=='Windows'", # 3.9.1及之后不提供windows的预编译包，避免一些没有编译环境的windows设备安装失败 + "matplotlib;platform_system=='Linux' or platform_system=='Darwin'", # linux 和 macos 不应限制matplotlib的最高版本，以避免无法更新导致的一些bug + "ultralytics", # yolov8,公式检测 + "paddleocr==2.7.3", # 2.8.0及2.8.1版本与detectron2有冲突，需锁定2.7.3 + "paddlepaddle==3.0.0b1;platform_system=='Linux'", # 解决linux的段异常问题 + "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", # windows版本3.0.0b1效率下降，需锁定2.6.1 + "pypandoc", # 表格解析latex转html + "struct-eqtable==0.1.0", # 表格解析 "detectron2" ], }, @@ -52,7 +55,8 @@ if __name__ == '__main__': python_requires=">=3.9", # 项目依赖的 Python 版本 entry_points={ "console_scripts": [ - "magic-pdf = magic_pdf.cli.magicpdf:cli" + "magic-pdf = magic_pdf.tools.cli:cli", + "magic-pdf-dev = magic_pdf.tools.cli_dev:cli" ], }, # 项目提供的可执行命令 include_package_data=True, # 是否包含非代码文件，如数据文件、配置文件等 diff --git a/magic_pdf/train_utils/__init__.py b/tests/test_tools/__init__.py similarity index 100% rename from magic_pdf/train_utils/__init__.py rename to tests/test_tools/__init__.py diff --git a/tests/test_tools/assets/cli/path/cli_test_01.pdf b/tests/test_tools/assets/cli/path/cli_test_01.pdf new file mode 100644 index 00000000..229be9ce Binary files /dev/null and b/tests/test_tools/assets/cli/path/cli_test_01.pdf differ diff --git a/tests/test_tools/assets/cli/path/cli_test_02.pdf b/tests/test_tools/assets/cli/path/cli_test_02.pdf new file mode 100644 index 00000000..1adcc01c Binary files /dev/null and b/tests/test_tools/assets/cli/path/cli_test_02.pdf differ diff --git a/tests/test_tools/assets/cli/pdf/cli_test_01.pdf b/tests/test_tools/assets/cli/pdf/cli_test_01.pdf new file mode 100644 index 00000000..229be9ce Binary files /dev/null and b/tests/test_tools/assets/cli/pdf/cli_test_01.pdf differ diff --git a/tests/test_tools/assets/cli_dev/cli_test_01.jsonl b/tests/test_tools/assets/cli_dev/cli_test_01.jsonl new file mode 100644 index 00000000..c9ca39c7 --- /dev/null +++ b/tests/test_tools/assets/cli_dev/cli_test_01.jsonl @@ -0,0 +1 @@ +{"file_location":"tests/test_tools/assets/cli_dev/cli_test_01.pdf","doc_layout_result":[{"layout_dets":[{"category_id":1,"poly":[882.4013061523438,169.93817138671875,1552.350341796875,169.93817138671875,1552.350341796875,625.8263549804688,882.4013061523438,625.8263549804688],"score":0.999992311000824},{"category_id":1,"poly":[882.474853515625,1450.92822265625,1551.4490966796875,1450.92822265625,1551.4490966796875,1877.5712890625,882.474853515625,1877.5712890625],"score":0.9999903440475464},{"category_id":1,"poly":[881.6513061523438,626.2058715820312,1552.1400146484375,626.2058715820312,1552.1400146484375,1450.604736328125,881.6513061523438,1450.604736328125],"score":0.9999856352806091},{"category_id":1,"poly":[149.41075134277344,232.1595001220703,819.0465087890625,232.1595001220703,819.0465087890625,625.8865356445312,149.41075134277344,625.8865356445312],"score":0.99998539686203},{"category_id":1,"poly":[149.3945770263672,1215.5172119140625,817.8850708007812,1215.5172119140625,817.8850708007812,1304.873291015625,149.3945770263672,1304.873291015625],"score":0.9999765157699585},{"category_id":1,"poly":[882.6979370117188,1880.13916015625,1552.15185546875,1880.13916015625,1552.15185546875,2031.339599609375,882.6979370117188,2031.339599609375],"score":0.9999744892120361},{"category_id":1,"poly":[148.96054077148438,743.3055419921875,818.6231689453125,743.3055419921875,818.6231689453125,1074.2369384765625,148.96054077148438,1074.2369384765625],"score":0.9999669790267944},{"category_id":1,"poly":[148.8435516357422,1791.14306640625,818.6885375976562,1791.14306640625,818.6885375976562,2030.794189453125,148.8435516357422,2030.794189453125],"score":0.9999618530273438},{"category_id":0,"poly":[150.7009735107422,684.0087890625,623.5106201171875,684.0087890625,623.5106201171875,717.03662109375,150.7009735107422,717.03662109375],"score":0.9999415278434753},{"category_id":8,"poly":[146.48068237304688,1331.6737060546875,317.2640075683594,1331.6737060546875,317.2640075683594,1400.1722412109375,146.48068237304688,1400.1722412109375],"score":0.9998958110809326},{"category_id":1,"poly":[149.42420959472656,1430.8782958984375,818.9042358398438,1430.8782958984375,818.9042358398438,1672.7386474609375,149.42420959472656,1672.7386474609375],"score":0.9998599290847778},{"category_id":1,"poly":[149.18746948242188,172.10252380371094,818.5662231445312,172.10252380371094,818.5662231445312,230.4594268798828,149.18746948242188,230.4594268798828],"score":0.9997718334197998},{"category_id":0,"poly":[149.0175018310547,1732.1090087890625,702.1005859375,1732.1090087890625,702.1005859375,1763.6046142578125,149.0175018310547,1763.6046142578125],"score":0.9997085928916931},{"category_id":2,"poly":[1519.802490234375,98.59099578857422,1551.985107421875,98.59099578857422,1551.985107421875,119.48420715332031,1519.802490234375,119.48420715332031],"score":0.9995552897453308},{"category_id":8,"poly":[146.9109649658203,1100.156494140625,544.2803344726562,1100.156494140625,544.2803344726562,1184.929443359375,146.9109649658203,1184.929443359375],"score":0.9995207786560059},{"category_id":2,"poly":[148.11611938476562,99.87767791748047,318.926025390625,99.87767791748047,318.926025390625,120.70393371582031,148.11611938476562,120.70393371582031],"score":0.999351441860199},{"category_id":9,"poly":[791.7642211914062,1130.056396484375,818.6940307617188,1130.056396484375,818.6940307617188,1161.1080322265625,791.7642211914062,1161.1080322265625],"score":0.9908884763717651},{"category_id":9,"poly":[788.37060546875,1346.8450927734375,818.5010986328125,1346.8450927734375,818.5010986328125,1377.370361328125,788.37060546875,1377.370361328125],"score":0.9873985052108765},{"category_id":14,"poly":[146,1103,543,1103,543,1184,146,1184],"score":0.94,"latex":"E\\!\\left(W\\right)\\!=\\!\\frac{E\\!\\left[H^{2}\\right]}{2E\\!\\left[H\\right]}\\!=\\!\\frac{E\\!\\left[H\\right]}{2}\\!\\!\\left(1\\!+\\!\\operatorname{CV}\\!\\left(H\\right)^{2}\\right)"},{"category_id":13,"poly":[1196,354,1278,354,1278,384,1196,384],"score":0.91,"latex":"p(1-q)"},{"category_id":13,"poly":[881,415,1020,415,1020,444,881,444],"score":0.91,"latex":"(1-p)(1-q)"},{"category_id":14,"poly":[147,1333,318,1333,318,1400,147,1400],"score":0.91,"latex":"\\mathbf{CV}\\big(H\\big)\\!=\\!\\frac{\\boldsymbol{\\upsigma}_{H}}{E\\big[H\\big]}"},{"category_id":13,"poly":[1197,657,1263,657,1263,686,1197,686],"score":0.9,"latex":"(1-p)"},{"category_id":13,"poly":[213,1217,263,1217,263,1244,213,1244],"score":0.88,"latex":"E[X]"},{"category_id":13,"poly":[214,1434,245,1434,245,1459,214,1459],"score":0.87,"latex":"\\upsigma_{H}"},{"category_id":13,"poly":[324,2002,373,2002,373,2028,324,2028],"score":0.84,"latex":"30\\%"},{"category_id":13,"poly":[1209,693,1225,693,1225,717,1209,717],"score":0.83,"latex":"p"},{"category_id":13,"poly":[990,449,1007,449,1007,474,990,474],"score":0.81,"latex":"p"},{"category_id":13,"poly":[346,1277,369,1277,369,1301,346,1301],"score":0.81,"latex":"H"},{"category_id":13,"poly":[1137,661,1154,661,1154,686,1137,686],"score":0.81,"latex":"p"},{"category_id":13,"poly":[522,1432,579,1432,579,1459,522,1459],"score":0.81,"latex":"H\\left(4\\right)"},{"category_id":13,"poly":[944,540,962,540,962,565,944,565],"score":0.8,"latex":"p"},{"category_id":13,"poly":[1444,936,1461,936,1461,961,1444,961],"score":0.79,"latex":"p"},{"category_id":13,"poly":[602,1247,624,1247,624,1270,602,1270],"score":0.78,"latex":"H"},{"category_id":13,"poly":[147,1247,167,1247,167,1271,147,1271],"score":0.77,"latex":"X"},{"category_id":13,"poly":[210,1246,282,1246,282,1274,210,1274],"score":0.77,"latex":"\\operatorname{CV}(H)"},{"category_id":13,"poly":[1346,268,1361,268,1361,292,1346,292],"score":0.76,"latex":"q"},{"category_id":13,"poly":[215,957,238,957,238,981,215,981],"score":0.74,"latex":"H"},{"category_id":13,"poly":[149,956,173,956,173,981,149,981],"score":0.63,"latex":"W"},{"category_id":13,"poly":[924,841,1016,841,1016,868,924,868],"score":0.56,"latex":"8{\\cdot}00\\;\\mathrm{a.m}"},{"category_id":13,"poly":[956,871,1032,871,1032,898,956,898],"score":0.43,"latex":"20~\\mathrm{min}"},{"category_id":13,"poly":[1082,781,1112,781,1112,808,1082,808],"score":0.41,"latex":"(l)"},{"category_id":13,"poly":[697,1821,734,1821,734,1847,697,1847],"score":0.3,"latex":"^{1\\mathrm{~h~}}"}],"page_info":{"page_no":0,"height":2200,"width":1700}}]} \ No newline at end of file diff --git a/tests/test_tools/assets/cli_dev/cli_test_01.model.json b/tests/test_tools/assets/cli_dev/cli_test_01.model.json new file mode 100644 index 00000000..a55f91fe --- /dev/null +++ b/tests/test_tools/assets/cli_dev/cli_test_01.model.json @@ -0,0 +1,638 @@ +[ + { + "layout_dets": [ + { + "category_id": 1, + "poly": [ + 882.4013061523438, + 169.93817138671875, + 1552.350341796875, + 169.93817138671875, + 1552.350341796875, + 625.8263549804688, + 882.4013061523438, + 625.8263549804688 + ], + "score": 0.999992311000824 + }, + { + "category_id": 1, + "poly": [ + 882.474853515625, + 1450.92822265625, + 1551.4490966796875, + 1450.92822265625, + 1551.4490966796875, + 1877.5712890625, + 882.474853515625, + 1877.5712890625 + ], + "score": 0.9999903440475464 + }, + { + "category_id": 1, + "poly": [ + 881.6513061523438, + 626.2058715820312, + 1552.1400146484375, + 626.2058715820312, + 1552.1400146484375, + 1450.604736328125, + 881.6513061523438, + 1450.604736328125 + ], + "score": 0.9999856352806091 + }, + { + "category_id": 1, + "poly": [ + 149.41075134277344, + 232.1595001220703, + 819.0465087890625, + 232.1595001220703, + 819.0465087890625, + 625.8865356445312, + 149.41075134277344, + 625.8865356445312 + ], + "score": 0.99998539686203 + }, + { + "category_id": 1, + "poly": [ + 149.3945770263672, + 1215.5172119140625, + 817.8850708007812, + 1215.5172119140625, + 817.8850708007812, + 1304.873291015625, + 149.3945770263672, + 1304.873291015625 + ], + "score": 0.9999765157699585 + }, + { + "category_id": 1, + "poly": [ + 882.6979370117188, + 1880.13916015625, + 1552.15185546875, + 1880.13916015625, + 1552.15185546875, + 2031.339599609375, + 882.6979370117188, + 2031.339599609375 + ], + "score": 0.9999744892120361 + }, + { + "category_id": 1, + "poly": [ + 148.96054077148438, + 743.3055419921875, + 818.6231689453125, + 743.3055419921875, + 818.6231689453125, + 1074.2369384765625, + 148.96054077148438, + 1074.2369384765625 + ], + "score": 0.9999669790267944 + }, + { + "category_id": 1, + "poly": [ + 148.8435516357422, + 1791.14306640625, + 818.6885375976562, + 1791.14306640625, + 818.6885375976562, + 2030.794189453125, + 148.8435516357422, + 2030.794189453125 + ], + "score": 0.9999618530273438 + }, + { + "category_id": 0, + "poly": [ + 150.7009735107422, + 684.0087890625, + 623.5106201171875, + 684.0087890625, + 623.5106201171875, + 717.03662109375, + 150.7009735107422, + 717.03662109375 + ], + "score": 0.9999415278434753 + }, + { + "category_id": 8, + "poly": [ + 146.48068237304688, + 1331.6737060546875, + 317.2640075683594, + 1331.6737060546875, + 317.2640075683594, + 1400.1722412109375, + 146.48068237304688, + 1400.1722412109375 + ], + "score": 0.9998958110809326 + }, + { + "category_id": 1, + "poly": [ + 149.42420959472656, + 1430.8782958984375, + 818.9042358398438, + 1430.8782958984375, + 818.9042358398438, + 1672.7386474609375, + 149.42420959472656, + 1672.7386474609375 + ], + "score": 0.9998599290847778 + }, + { + "category_id": 1, + "poly": [ + 149.18746948242188, + 172.10252380371094, + 818.5662231445312, + 172.10252380371094, + 818.5662231445312, + 230.4594268798828, + 149.18746948242188, + 230.4594268798828 + ], + "score": 0.9997718334197998 + }, + { + "category_id": 0, + "poly": [ + 149.0175018310547, + 1732.1090087890625, + 702.1005859375, + 1732.1090087890625, + 702.1005859375, + 1763.6046142578125, + 149.0175018310547, + 1763.6046142578125 + ], + "score": 0.9997085928916931 + }, + { + "category_id": 2, + "poly": [ + 1519.802490234375, + 98.59099578857422, + 1551.985107421875, + 98.59099578857422, + 1551.985107421875, + 119.48420715332031, + 1519.802490234375, + 119.48420715332031 + ], + "score": 0.9995552897453308 + }, + { + "category_id": 8, + "poly": [ + 146.9109649658203, + 1100.156494140625, + 544.2803344726562, + 1100.156494140625, + 544.2803344726562, + 1184.929443359375, + 146.9109649658203, + 1184.929443359375 + ], + "score": 0.9995207786560059 + }, + { + "category_id": 2, + "poly": [ + 148.11611938476562, + 99.87767791748047, + 318.926025390625, + 99.87767791748047, + 318.926025390625, + 120.70393371582031, + 148.11611938476562, + 120.70393371582031 + ], + "score": 0.999351441860199 + }, + { + "category_id": 9, + "poly": [ + 791.7642211914062, + 1130.056396484375, + 818.6940307617188, + 1130.056396484375, + 818.6940307617188, + 1161.1080322265625, + 791.7642211914062, + 1161.1080322265625 + ], + "score": 0.9908884763717651 + }, + { + "category_id": 9, + "poly": [ + 788.37060546875, + 1346.8450927734375, + 818.5010986328125, + 1346.8450927734375, + 818.5010986328125, + 1377.370361328125, + 788.37060546875, + 1377.370361328125 + ], + "score": 0.9873985052108765 + }, + { + "category_id": 14, + "poly": [ + 146, + 1103, + 543, + 1103, + 543, + 1184, + 146, + 1184 + ], + "score": 0.94, + "latex": "E\\!\\left(W\\right)\\!=\\!\\frac{E\\!\\left[H^{2}\\right]}{2E\\!\\left[H\\right]}\\!=\\!\\frac{E\\!\\left[H\\right]}{2}\\!\\!\\left(1\\!+\\!\\operatorname{CV}\\!\\left(H\\right)^{2}\\right)" + }, + { + "category_id": 13, + "poly": [ + 1196, + 354, + 1278, + 354, + 1278, + 384, + 1196, + 384 + ], + "score": 0.91, + "latex": "p(1-q)" + }, + { + "category_id": 13, + "poly": [ + 881, + 415, + 1020, + 415, + 1020, + 444, + 881, + 444 + ], + "score": 0.91, + "latex": "(1-p)(1-q)" + }, + { + "category_id": 14, + "poly": [ + 147, + 1333, + 318, + 1333, + 318, + 1400, + 147, + 1400 + ], + "score": 0.91, + "latex": "\\mathbf{CV}\\big(H\\big)\\!=\\!\\frac{\\boldsymbol{\\upsigma}_{H}}{E\\big[H\\big]}" + }, + { + "category_id": 13, + "poly": [ + 1197, + 657, + 1263, + 657, + 1263, + 686, + 1197, + 686 + ], + "score": 0.9, + "latex": "(1-p)" + }, + { + "category_id": 13, + "poly": [ + 213, + 1217, + 263, + 1217, + 263, + 1244, + 213, + 1244 + ], + "score": 0.88, + "latex": "E[X]" + }, + { + "category_id": 13, + "poly": [ + 214, + 1434, + 245, + 1434, + 245, + 1459, + 214, + 1459 + ], + "score": 0.87, + "latex": "\\upsigma_{H}" + }, + { + "category_id": 13, + "poly": [ + 324, + 2002, + 373, + 2002, + 373, + 2028, + 324, + 2028 + ], + "score": 0.84, + "latex": "30\\%" + }, + { + "category_id": 13, + "poly": [ + 1209, + 693, + 1225, + 693, + 1225, + 717, + 1209, + 717 + ], + "score": 0.83, + "latex": "p" + }, + { + "category_id": 13, + "poly": [ + 990, + 449, + 1007, + 449, + 1007, + 474, + 990, + 474 + ], + "score": 0.81, + "latex": "p" + }, + { + "category_id": 13, + "poly": [ + 346, + 1277, + 369, + 1277, + 369, + 1301, + 346, + 1301 + ], + "score": 0.81, + "latex": "H" + }, + { + "category_id": 13, + "poly": [ + 1137, + 661, + 1154, + 661, + 1154, + 686, + 1137, + 686 + ], + "score": 0.81, + "latex": "p" + }, + { + "category_id": 13, + "poly": [ + 522, + 1432, + 579, + 1432, + 579, + 1459, + 522, + 1459 + ], + "score": 0.81, + "latex": "H\\left(4\\right)" + }, + { + "category_id": 13, + "poly": [ + 944, + 540, + 962, + 540, + 962, + 565, + 944, + 565 + ], + "score": 0.8, + "latex": "p" + }, + { + "category_id": 13, + "poly": [ + 1444, + 936, + 1461, + 936, + 1461, + 961, + 1444, + 961 + ], + "score": 0.79, + "latex": "p" + }, + { + "category_id": 13, + "poly": [ + 602, + 1247, + 624, + 1247, + 624, + 1270, + 602, + 1270 + ], + "score": 0.78, + "latex": "H" + }, + { + "category_id": 13, + "poly": [ + 147, + 1247, + 167, + 1247, + 167, + 1271, + 147, + 1271 + ], + "score": 0.77, + "latex": "X" + }, + { + "category_id": 13, + "poly": [ + 210, + 1246, + 282, + 1246, + 282, + 1274, + 210, + 1274 + ], + "score": 0.77, + "latex": "\\operatorname{CV}(H)" + }, + { + "category_id": 13, + "poly": [ + 1346, + 268, + 1361, + 268, + 1361, + 292, + 1346, + 292 + ], + "score": 0.76, + "latex": "q" + }, + { + "category_id": 13, + "poly": [ + 215, + 957, + 238, + 957, + 238, + 981, + 215, + 981 + ], + "score": 0.74, + "latex": "H" + }, + { + "category_id": 13, + "poly": [ + 149, + 956, + 173, + 956, + 173, + 981, + 149, + 981 + ], + "score": 0.63, + "latex": "W" + }, + { + "category_id": 13, + "poly": [ + 924, + 841, + 1016, + 841, + 1016, + 868, + 924, + 868 + ], + "score": 0.56, + "latex": "8{\\cdot}00\\;\\mathrm{a.m}" + }, + { + "category_id": 13, + "poly": [ + 956, + 871, + 1032, + 871, + 1032, + 898, + 956, + 898 + ], + "score": 0.43, + "latex": "20~\\mathrm{min}" + }, + { + "category_id": 13, + "poly": [ + 1082, + 781, + 1112, + 781, + 1112, + 808, + 1082, + 808 + ], + "score": 0.41, + "latex": "(l)" + }, + { + "category_id": 13, + "poly": [ + 697, + 1821, + 734, + 1821, + 734, + 1847, + 697, + 1847 + ], + "score": 0.3, + "latex": "^{1\\mathrm{~h~}}" + } + ], + "page_info": { + "page_no": 0, + "height": 2200, + "width": 1700 + } + } +] \ No newline at end of file diff --git a/tests/test_tools/assets/cli_dev/cli_test_01.pdf b/tests/test_tools/assets/cli_dev/cli_test_01.pdf new file mode 100644 index 00000000..229be9ce Binary files /dev/null and b/tests/test_tools/assets/cli_dev/cli_test_01.pdf differ diff --git a/tests/test_tools/assets/common/cli_test_01.pdf b/tests/test_tools/assets/common/cli_test_01.pdf new file mode 100644 index 00000000..229be9ce Binary files /dev/null and b/tests/test_tools/assets/common/cli_test_01.pdf differ diff --git a/tests/test_tools/test_cli.py b/tests/test_tools/test_cli.py new file mode 100644 index 00000000..e66b6d63 --- /dev/null +++ b/tests/test_tools/test_cli.py @@ -0,0 +1,125 @@ +import tempfile +import os +import shutil +from click.testing import CliRunner + +from magic_pdf.tools.cli import cli + + +def test_cli_pdf(): + # setup + unitest_dir = "/tmp/magic_pdf/unittest/tools" + filename = "cli_test_01" + os.makedirs(unitest_dir, exist_ok=True) + temp_output_dir = tempfile.mkdtemp(dir="/tmp/magic_pdf/unittest/tools") + + # run + runner = CliRunner() + result = runner.invoke( + cli, + [ + "-p", + "tests/test_tools/assets/cli/pdf/cli_test_01.pdf", + "-o", + temp_output_dir, + ], + ) + + # check + assert result.exit_code == 0 + + base_output_dir = os.path.join(temp_output_dir, "cli_test_01/auto") + + r = os.stat(os.path.join(base_output_dir, f"{filename}.md")) + assert r.st_size > 7000 + + r = os.stat(os.path.join(base_output_dir, "middle.json")) + assert r.st_size > 200000 + + r = os.stat(os.path.join(base_output_dir, "model.json")) + assert r.st_size > 15000 + + r = os.stat(os.path.join(base_output_dir, "origin.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "layout.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "spans.pdf")) + assert r.st_size > 500000 + + assert os.path.exists(os.path.join(base_output_dir, "images")) is True + assert os.path.isdir(os.path.join(base_output_dir, "images")) is True + assert os.path.exists(os.path.join(base_output_dir, "content_list.json")) is False + + # teardown + shutil.rmtree(temp_output_dir) + + +def test_cli_path(): + # setup + unitest_dir = "/tmp/magic_pdf/unittest/tools" + os.makedirs(unitest_dir, exist_ok=True) + temp_output_dir = tempfile.mkdtemp(dir="/tmp/magic_pdf/unittest/tools") + + # run + runner = CliRunner() + result = runner.invoke( + cli, ["-p", "tests/test_tools/assets/cli/path", "-o", temp_output_dir] + ) + + # check + assert result.exit_code == 0 + + filename = "cli_test_01" + base_output_dir = os.path.join(temp_output_dir, "cli_test_01/auto") + + r = os.stat(os.path.join(base_output_dir, f"{filename}.md")) + assert r.st_size > 7000 + + r = os.stat(os.path.join(base_output_dir, "middle.json")) + assert r.st_size > 200000 + + r = os.stat(os.path.join(base_output_dir, "model.json")) + assert r.st_size > 15000 + + r = os.stat(os.path.join(base_output_dir, "origin.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "layout.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "spans.pdf")) + assert r.st_size > 500000 + + assert os.path.exists(os.path.join(base_output_dir, "images")) is True + assert os.path.isdir(os.path.join(base_output_dir, "images")) is True + assert os.path.exists(os.path.join(base_output_dir, "content_list.json")) is False + + base_output_dir = os.path.join(temp_output_dir, "cli_test_02/auto") + filename = "cli_test_02" + + r = os.stat(os.path.join(base_output_dir, f"{filename}.md")) + assert r.st_size > 5000 + + r = os.stat(os.path.join(base_output_dir, "middle.json")) + assert r.st_size > 200000 + + r = os.stat(os.path.join(base_output_dir, "model.json")) + assert r.st_size > 15000 + + r = os.stat(os.path.join(base_output_dir, "origin.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "layout.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "spans.pdf")) + assert r.st_size > 500000 + + assert os.path.exists(os.path.join(base_output_dir, "images")) is True + assert os.path.isdir(os.path.join(base_output_dir, "images")) is True + assert os.path.exists(os.path.join(base_output_dir, "content_list.json")) is False + + # teardown + shutil.rmtree(temp_output_dir) diff --git a/tests/test_tools/test_cli_dev.py b/tests/test_tools/test_cli_dev.py new file mode 100644 index 00000000..56470237 --- /dev/null +++ b/tests/test_tools/test_cli_dev.py @@ -0,0 +1,120 @@ +import tempfile +import os +import shutil +from click.testing import CliRunner + +from magic_pdf.tools import cli_dev + + +def test_cli_pdf(): + # setup + unitest_dir = "/tmp/magic_pdf/unittest/tools" + filename = "cli_test_01" + os.makedirs(unitest_dir, exist_ok=True) + temp_output_dir = tempfile.mkdtemp(dir="/tmp/magic_pdf/unittest/tools") + + # run + runner = CliRunner() + result = runner.invoke( + cli_dev.cli, + [ + "pdf", + "-p", + "tests/test_tools/assets/cli/pdf/cli_test_01.pdf", + "-j", + "tests/test_tools/assets/cli_dev/cli_test_01.model.json", + "-o", + temp_output_dir, + ], + ) + + # check + assert result.exit_code == 0 + + base_output_dir = os.path.join(temp_output_dir, "cli_test_01/auto") + + r = os.stat(os.path.join(base_output_dir, "content_list.json")) + assert r.st_size > 5000 + + r = os.stat(os.path.join(base_output_dir, f"{filename}.md")) + assert r.st_size > 7000 + + r = os.stat(os.path.join(base_output_dir, "middle.json")) + assert r.st_size > 200000 + + r = os.stat(os.path.join(base_output_dir, "model.json")) + assert r.st_size > 15000 + + r = os.stat(os.path.join(base_output_dir, "origin.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "layout.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "spans.pdf")) + assert r.st_size > 500000 + + assert os.path.exists(os.path.join(base_output_dir, "images")) is True + assert os.path.isdir(os.path.join(base_output_dir, "images")) is True + + # teardown + shutil.rmtree(temp_output_dir) + + +def test_cli_jsonl(): + # setup + unitest_dir = "/tmp/magic_pdf/unittest/tools" + filename = "cli_test_01" + os.makedirs(unitest_dir, exist_ok=True) + temp_output_dir = tempfile.mkdtemp(dir="/tmp/magic_pdf/unittest/tools") + + def mock_read_s3_path(s3path): + with open(s3path, "rb") as f: + return f.read() + + cli_dev.read_s3_path = mock_read_s3_path # mock + + # run + runner = CliRunner() + result = runner.invoke( + cli_dev.cli, + [ + "jsonl", + "-j", + "tests/test_tools/assets/cli_dev/cli_test_01.jsonl", + "-o", + temp_output_dir, + ], + ) + + # check + assert result.exit_code == 0 + + base_output_dir = os.path.join(temp_output_dir, "cli_test_01/auto") + + r = os.stat(os.path.join(base_output_dir, "content_list.json")) + assert r.st_size > 5000 + + r = os.stat(os.path.join(base_output_dir, f"{filename}.md")) + assert r.st_size > 7000 + + r = os.stat(os.path.join(base_output_dir, "middle.json")) + assert r.st_size > 200000 + + r = os.stat(os.path.join(base_output_dir, "model.json")) + assert r.st_size > 15000 + + r = os.stat(os.path.join(base_output_dir, "origin.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "layout.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "spans.pdf")) + assert r.st_size > 500000 + + assert os.path.exists(os.path.join(base_output_dir, "images")) is True + assert os.path.isdir(os.path.join(base_output_dir, "images")) is True + + # teardown + shutil.rmtree(temp_output_dir) diff --git a/tests/test_tools/test_common.py b/tests/test_tools/test_common.py new file mode 100644 index 00000000..0d38e1de --- /dev/null +++ b/tests/test_tools/test_common.py @@ -0,0 +1,52 @@ +import tempfile +import os +import shutil + +import pytest + +from magic_pdf.tools.common import do_parse + + +@pytest.mark.parametrize("method", ["auto", "txt", "ocr"]) +def test_common_do_parse(method): + # setup + unitest_dir = "/tmp/magic_pdf/unittest/tools" + filename = "fake" + os.makedirs(unitest_dir, exist_ok=True) + + temp_output_dir = tempfile.mkdtemp(dir="/tmp/magic_pdf/unittest/tools") + + # run + with open("tests/test_tools/assets/common/cli_test_01.pdf", "rb") as f: + bits = f.read() + do_parse(temp_output_dir, filename, bits, [], method, f_dump_content_list=True) + + # check + base_output_dir = os.path.join(temp_output_dir, f"fake/{method}") + + r = os.stat(os.path.join(base_output_dir, "content_list.json")) + assert r.st_size > 5000 + + r = os.stat(os.path.join(base_output_dir, f"{filename}.md")) + assert r.st_size > 7000 + + r = os.stat(os.path.join(base_output_dir, "middle.json")) + assert r.st_size > 200000 + + r = os.stat(os.path.join(base_output_dir, "model.json")) + assert r.st_size > 15000 + + r = os.stat(os.path.join(base_output_dir, "origin.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "layout.pdf")) + assert r.st_size > 500000 + + r = os.stat(os.path.join(base_output_dir, "spans.pdf")) + assert r.st_size > 500000 + + os.path.exists(os.path.join(base_output_dir, "images")) + os.path.isdir(os.path.join(base_output_dir, "images")) + + # teardown + shutil.rmtree(temp_output_dir)