【文档智能】实践：基于Yolo三行代码极简的训练一个版式分析模型

一、数据集

本文以开源的CDLA数据集做为实验，CDLA是一个中文文档版面分析数据集，面向中文文献类（论文）场景。包含以下10个label：

数据集下载地址：https://github.com/buptlihang/CDLA

数据集是labelme格式的文件，需要自行转化成yolo训练所需要的格式，最后文件放置形式如下：

注意：这个数据集有个问题就是缺乏段落标注，含相关段落标注的开源模型可以查看上篇 《【文档智能】包含段落的开源的中文版面分析模型》

二、安装所需要的依赖

        
          
pip install ultralytics  

      

三、配置yaml文件

cdla.yaml

        
          
#path: ./datasets/  # dataset root dir  
train: cdla/images/train  
val: cdla/images/dev  
test:  # test images (optional)  
  
nc: 10   # 数据集类别数量  
names: [ # 数据集标签  
  'Text',  
  'Title',  
  'Figure',  
  'Figure caption',  
  'Table',  
  'Table caption',  
  'Header',  
  'Footer',  
  'Reference',  
  'Equation',  
]  
  

      

三、训练代码

        
          
from ultralytics import YOLO  
  
# model = YOLO('yolov8n.yaml')    # 重新训练模型  
model = YOLO('yolov8n.pt')        # 加载预训练模型  
  
# Train the model  
results = model.train(data='cdla.yaml', epochs=100, imgsz=640, device=0)  

      

四、预测代码

        
          
from ultralytics import YOLO  
  
image_path = ''  # 待预测图片路径  
model_path = ''  # 权重路径  
model = YOLO(model_path)  
  
result = model(image_path, save=True, conf=0.5, save_crop=False, line_width=2)  
print(result)  
  
print(result[0].names)         # 输出id2label map  
print(result[0].boxes)         # 输出所有的检测到的bounding box  
print(result[0].boxes.xyxy)    # 输出所有的检测到的bounding box的左上和右下坐标  
print(result[0].boxes.cls)     # 输出所有的检测到的bounding box类别对应的id  
print(result[0].boxes.conf)    # 输出所有的检测到的bounding box的置信度  

      

往期相关

【文档智能 & RAG】RAG增强之路：增强PDF解析并结构化技术路线方案及思路

【文档智能 & LLM】LayoutLLM：一种多模态文档布局模型和大模型结合的框架

【文档智能】再谈基于Transformer架构的文档智能理解方法论和相关数据集

【文档智能】多模态预训练模型及相关数据集汇总

【文档智能】：GeoLayoutLM：一种用于视觉信息提取（VIE）的多模态预训练模型

文档智能：ERNIE-Layout

【文档智能】符合人类阅读顺序的文档模型-LayoutReader及非官方权重开源