摘要

我们提出DocLayout-YOLO，基于YOLO-v10构建了一个实时且鲁棒的文档布局检测模型。该模型通过多样化文档预训练和针对布局检测的结构优化进行增强。在预训练阶段，我们引入Mesh-candidate BestFit，将文档合成视为二维装箱问题，并创建了大规模多样化合成文档数据集DocSynth-300K。在模型结构优化方面，我们提出了一种具有全局到局部可控性的模块，用于精确检测不同尺度下的文档元素。

新闻 🚀🚀🚀

2024.10.25 🎉🎉 Mesh-candidate Bestfit 代码已发布。Mesh-candidate Bestfit 是一个自动流程，可生成大规模、高质量且视觉吸引人的文档布局检测数据集。教程和示例数据见 此处。

2024.10.23 🎉🎉 DocSynth300K数据集 已发布在 🤗Huggingface，DocSynth300K 是一个大规模且多样化的文档布局分析预训练数据集，可显著提升模型性能。

2024.10.21 🎉🎉 在线演示 可在 🤗Huggingface 上使用。

2024.10.18 🎉🎉 DocLayout-YOLO 已在 PDF-Extract-Kit 中实现，用于文档上下文提取。

2024.10.16 🎉🎉 论文 现在可在 ArXiv 上获取。

快速入门

在线演示 现已可用。对于本地开发，请按照以下步骤操作：

1. 环境设置

按照以下步骤设置环境：

conda create -n doclayout_yolo python=3.10
conda activate doclayout_yolo
pip install -e .

注意： 如果仅需用于推理，可通过pip直接安装：

pip install doclayout-yolo

2. 预测

可以使用脚本或SDK进行预测：

• 脚本

  运行以下命令使用脚本进行预测：

  python demo.py --model path/to/model --image-path path/to/image
  

• SDK

  以下是使用SDK进行预测的示例：

  import cv2
  from doclayout_yolo import YOLOv10
  
  # 加载预训练模型
  model = YOLOv10("path/to/provided/model")
  
  # 进行预测
  det_res = model.predict(
      "path/to/image",   # 要预测的图像
      imgsz=1024,        # 预测图像大小
      conf=0.2,          # 置信度阈值
      device="cuda:0"    # 使用的设备（例如，'cuda:0' 或 'cpu'）
  )
  
  # 注释并保存结果
  annotated_frame = det_res[0].plot(pil=True, line_width=5, font_size=20)
  cv2.imwrite("result.jpg", annotated_frame)
  

我们提供在DocStructBench上微调的模型用于预测，能够处理各种文档类型。模型可从 此处 下载，示例图片见 assets/example 目录。

注意： 对于PDF内容提取，请参考 PDF-Extract-Kit 和 MinerU。

注意： 感谢 NielsRogge，DocLayout-YOLO 现在支持直接从 🤗Huggingface 实现，可按以下方式加载模型：

filepath = hf_hub_download(repo_id="juliozhao/DocLayout-YOLO-DocStructBench", filename="doclayout_yolo_docstructbench_imgsz1024.pt")
model = YOLOv10(filepath)

或直接使用 from_pretrained：

model = YOLOv10.from_pretrained("juliozhao/DocLayout-YOLO-DocStructBench")

更多详情见 此PR。

注意： 感谢 luciaganlulu，DocLayout-YOLO 可进行批量推理和预测。不同于 demo.py 中单张图像传入 model.predict，应传入图像路径列表。此外，由于在 YOLOv11 之前未实现批量推理，需手动修改 此处 的 batch_size。

DocSynth300K 数据集

数据下载

使用以下命令下载数据集（约113G）：

from huggingface_hub import snapshot_download
# 下载 DocSynth300K
snapshot_download(repo_id="juliozhao/DocSynth300K", local_dir="./docsynth300k-hf", repo_type="dataset")
# 如果下载中断且文件不完整，可恢复下载
snapshot_download(repo_id="juliozhao/DocSynth300K", local_dir="./docsynth300k-hf", repo_type="dataset", resume_download=True)

数据格式化与预训练

若想进行 DocSynth300K 预训练，使用 format_docsynth300k.py 将原始 .parquet 格式转换为 YOLO 格式。转换后的数据将存储在 ./layout_data/docsynth300k。

python format_docsynth300k.py

进行 DocSynth300K 预训练，请使用此 命令。我们默认使用8块GPU进行预训练。为了达到最佳性能，可根据下游微调数据分布或设置调整超参数，如 imgsz、lr 等。

注意： 由于YOLO原始数据加载代码存在内存泄漏，大规模数据集的预训练可能意外中断，使用 --pretrain last_checkpoint.pt --resume 可恢复预训练过程。

在公共DLA数据集上的训练与评估

数据准备

1. 指定数据根路径

找到您的 ultralytics 配置文件（Linux 用户可在 $HOME/.config/Ultralytics/settings.yaml) 中查找）并修改 datasets_dir 为项目根路径。

2. 从以下链接下载准备好的 yolo 格式 D4LA 和 DocLayNet 数据并放入 ./layout_data：

数据集	下载
D4LA	链接
DocLayNet	链接

文件结构如下：

./layout_data
├── D4LA
│   ├── images
│   ├── labels
│   ├── test.txt
│   └── train.txt
└── doclaynet
    ├── images
    ├── labels
    ├── val.txt
    └── train.txt

训练与评估

训练使用 8 块 GPU，全局批量大小为 64（每块设备 8 张图像）。详细设置和检查点如下：

数据集	模型	DocSynth300K 预训练？	图像尺寸	学习率	微调	评估	AP50	mAP	检查点
D4LA	DocLayout-YOLO	❌	1600	0.04	命令	命令	81.7	69.8	检查点
D4LA	DocLayout-YOLO	✅	1600	0.04	命令	命令	82.4	70.3	检查点
DocLayNet	DocLayout-YOLO	❌	1120	0.02	命令	命令	93.0	77.7	检查点
DocLayNet	DocLayout-YOLO	✅	1120	0.02	命令	命令	93.4	79.7	检查点

DocSynth300K 预训练模型可从 此处 下载。评估时请将 checkpoint.pt 改为要评估的模型路径。

致谢

代码基于 ultralytics 和 YOLO-v10 开发。

感谢他们的出色工作！

星星历史

如果您认为该项目有用，请为仓库添加一个“星标”。看到您的关注让我们感到非常兴奋，这激励我们继续投入该项目！

引用

@misc{zhao2024doclayoutyoloenhancingdocumentlayout,
      title={DocLayout-YOLO: Enhancing Document Layout Analysis through Diverse Synthetic Data and Global-to-Local Adaptive Perception}, 
      author={Zhiyuan Zhao and Hengrui Kang and Bin Wang and Conghui He},
      year={2024},
      eprint={2410.12628},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2410.12628}, 
}

@article{wang2024mineru,
  title={MinerU: An Open-Source Solution for Precise Document Content Extraction},
  author={Wang, Bin and Xu, Chao and Zhao, Xiaomeng and Ouyang, Linke and Wu, Fan and Zhao, Zhiyuan and Xu, Rui and Liu, Kaiwen and Qu, Yuan and Shang, Fukai and others},
  journal={arXiv preprint arXiv:2409.18839},
  year={2024}
}

DocLayout-YOLO 快速上手指南

环境准备

• 系统要求：Python 3.10
• 前置依赖：
  • PyTorch
  • Ultralytics（YOLOv10框架）
  • Hugging Face Transformers

安装步骤

1. 创建虚拟环境并安装依赖：

conda create -n doclayout_yolo python=3.10
conda activate doclayout_yolo
pip install -e .

2. 仅需推理时可直接安装：

pip install doclayout-yolo

基本使用

1. 在线演示

访问 https://huggingface.co/spaces/opendatalab/DocLayout-YOLO 进行实时测试

2. 本地运行

使用脚本

python demo.py --model path/to/model --image-path path/to/image

使用SDK

import cv2
from doclayout_yolo import YOLOv10

model = YOLOv10("path/to/provided/model")
det_res = model.predict(
    "path/to/image", 
    imgsz=1024, 
    conf=0.2, 
    device="cuda:0"
)
annotated_frame = det_res[0].plot(pil=True, line_width=5, font_size=20)
cv2.imwrite("result.jpg", annotated_frame)

3. 模型获取

• 预训练模型：https://huggingface.co/juliozhao/DocLayout-YOLO-DocStructBench
• 示例图片：assets/example 目录

4. 数据集

• DocSynth300K：https://huggingface.co/datasets/juliozhao/DocSynth300K