PDF-Extract-Kit:迄今为止最好的PDF内容抽取开源项目

作者：lucas大叔

原文地址：https://zhuanlan.zhihu.com/p/709538411

Github 地址：https://github.com/opendatalab/PDF-Extract-Kit

整体介绍

PDF文档包含了大量的有用信息，例如文本、表格、图像、公式等。此外，PDF的文档布局也相当复杂，页眉、页脚、表格标题、图片标题等等。因此，实现高质量的PDF内容提取是一项非常有挑战性的工作。笔者也曾尝试过基于python包、基于布局检测+OCR识别的多种方案，始终没有找到满意的解决方案。

PDF-Extract-Kit 项目结合多样性PDF文档标注，训练了鲁棒的布局检测和公式检测模型。在论文、教材、研报、财报等多样性的PDF文档上，pipeline都能得到准确的提取结果，对于扫描模糊、水印等情况也有较高鲁棒性。

项目将PDF内容提取工作拆解为以下任务：

• 布局检测：使用LayoutLMv3模型进行区域检测，如图像，表格,标题,文本等；

• 公式检测：使用YOLOv8进行公式检测，包含行内公式和行间公式；

• 公式识别：使用UniMERNet进行公式识别；

• 光学字符识别：使用PaddleOCR进行文本识别；

PDF-Extract-Kit输出格式

{
    "layout_dets": [    # 页中的元素
        {
            "category_id": 0, # 类别编号， 0~9，13~15
            "poly": [
                136.0, # 坐标为图片坐标，需要转换回pdf坐标, 顺序是 左上-右上-右下-左下的x,y坐标
                781.0,
                340.0,
                781.0,
                340.0,
                806.0,
                136.0,
                806.0
            ],
            "score": 0.69,   # 置信度
            "latex": ''      # 公式识别的结果，只有13,14有内容，其他为空，另外15是ocr的结果，这个key会换成text
        },
        ...
    ],
    "page_info": {         # 页信息：提取bbox时的分辨率大小，如果有缩放可以基于该信息进行对齐
        "page_no": 0,      # 页数
        "height": 1684,    # 页高
        "width": 1200      # 页宽
    }
}

category_id类型的定义如下：

{0: 'title',              # 标题
 1: 'plain text',         # 文本
 2: 'abandon',            # 包括页眉页脚页码和页面注释
 3: 'figure',             # 图片
 4: 'figure_caption',     # 图片描述
 5: 'table',              # 表格
 6: 'table_caption',      # 表格描述
 7: 'table_footnote',     # 表格注释
 8: 'isolate_formula',    # 行间公式（这个是layout的行间公式，优先级低于14）
 9: 'formula_caption',    # 行间公式的标号

 13: 'inline_formula',    # 行内公式
 14: 'isolated_formula',  # 行间公式
 15: 'ocr_text'}              # ocr识别结果

效果展示

结合多样性PDF文档标注，我们训练了鲁棒的布局检测和公式检测模型。在论文、教材、研报、财报等多样性的PDF文档上，我们的pipeline都能得到准确的提取结果，对于扫描模糊、水印等情况也有较高鲁棒性。

实测的年报识别效果

从实测的结果可以看到，PDF-Extract-Kit可以识别出年报的表格描述，而其他的解决方案更多地是识别成普通的文本。美中不足的是，要想识别出表格的具体内容需要进一步的处理，OCR识别的文本也需要自己拼接。

评测指标

现有开源模型多基于Arxiv论文类型数据进行训练，面对多样性的PDF文档，提前质量远不能达到实用需求。相比之下，我们的模型经过多样化数据训练，可以适应各种类型文档提取。

评测代码及详细信息请看这里。

布局检测

我们与现有的开源Layout检测模型做了对比，包括DocXchain、Surya、360LayoutAnalysis的两个模型。而LayoutLMv3-SFT指的是我们在LayoutLMv3-base-chinese预训练权重的基础上进一步做了SFT训练后的模型。论文验证集由402张论文页面构成，教材验证集由587张不同来源的教材页面构成。

公式检测

我们与开源的模型Pix2Text-MFD做了对比。另外，YOLOv8-Trained是我们在YOLOv8l模型的基础上训练后的权重。论文验证集由255张论文页面构成，多源验证集由789张不同来源的页面构成，包括教材、书籍等。

公式识别

公式识别我们使用的是UniMERNet的权重，没有进一步的SFT训练，其精度验证结果可以在其GitHub页面获取。

使用教程

环境安装 (Linux)

在Windows上运行，请参考在Windows环境下使用PDF-Extract-Kit。

在macOS上运行，请参考在macOS系统使用PDF-Extract-Kit。

conda create -n pipeline python=3.10

pip install -r requirements.txt

pip install --extra-index-url https://miropsota.github.io/torch_packages_builder detectron2==0.6+pt2.3.1cu121

遇到pillow相关的报错，请安装下面指定的版本

pip install pillow==8.4.0

运行提取脚本

python pdf_extract.py --pdf data/pdfs/ocr_1.pdf

有四个参数：

• --pdf 待处理的pdf文件，如果传入一个文件夹，则会处理文件夹下的所有pdf文件。

• --output 处理结果保存的路径，默认是"output"

• --vis 是否对结果可视化，是则会把检测的结果可视化出来，主要是检测框和类别

• --render 是否把识别得的结果渲染出来，包括公式的latex代码，以及普通文本，都会渲染出来放在检测框中。注意：此过程非常耗时，另外也需要提前安装xelatex和imagemagic。