基于CDLA版面分析数据集使用Yolov8进行文档版面分析实战
1 yolov8
github：https://github.com/ultralytics/ultralytics
YOLOv8是一种先进的目标检测算法，它属于You Only Look Once（YOLO）系列的第八代版本。YOLO系列算法因其速度快、性能好而广受欢迎，尤其适用于实时目标检测任务。YOLOv8在继承前代优点的基础上，进一步优化了模型结构和检测性能，提高了检测速度和准确性，同时在处理复杂场景和小型目标上也有显著提升。
这篇文章我们就是用yolov8在中文版面分析数据集上做一个小小实践。
2 CDLA数据集
2.1 CDLA数据集
github：https://github.com/buptlihang/CDLA
CDLA是一个由北京邮电大学李航教授团队创建的中文文档版面分析（Chinese Document Layout Analysis）数据集。该数据集旨在支持文档图像处理和文档分析领域的研究，特别关注中文文档。它包含了标注好的文档图像，其中包含了标题（Title）、文本（Text）和页脚（Footer）等不同区域的标注信息。
2.2 CDLA数据集下载
在CDLA github仓库主页提供了数据集的下载链接:
- 百度网盘：https://pan.baidu.com/s/1449mhds2ze5JLk-88yKVAA，提取码：tp0d
- Google Drive：https://drive.google.com/file/d/14SUsp_TG8OPdK0VthRXBcAbYzIBjSNLm/view?usp=sharing
从上述地址下载CDLA数据集到本地。
2.3 CDLA数据集转Yolo格式
CDLA是使用labelme标注的数据集，标注文件包含的内容如下：
- "shapes": shapes字段是一个list，里面有多个dict，每个dict代表一个标注实例。
- "labels": 类别。
- "points": 实例标注。因为我们的标注是Polygon形式，所以points里的坐标数量可能大于4。
- "shape_type": "polygon"
- "imagePath": 图片路径/名
- "imageHeight": 高
- "imageWidth": 宽
示例的标注文件内容如下：
{
"version":"4.5.6",
"flags":{},
"shapes":[
{
"label":"Title",
"points":[
[
553.1111111111111,
166.59259259259258
],
[
553.1111111111111,
198.59259259259258
],
[
686.1111111111111,
198.59259259259258
],
[
686.1111111111111,
166.59259259259258
]
],
"group_id":null,
"shape_type":"polygon",
"flags":{}
},
{
"label":"Text",
"points":[
[
250.5925925925925,
298.0740740740741
],
[
250.5925925925925,
345.0740740740741
],
[
188.5925925925925,
345.0740740740741
],
[
188.5925925925925,
410.0740740740741
],
[
188.5925925925925,
456.0740740740741
],
[
324.5925925925925,
456.0740740740741
],
[
324.5925925925925,
410.0740740740741
],
[
1051.5925925925926,
410.0740740740741
],
[
1051.5925925925926,
345.0740740740741
],
[
1052.5925925925926,
345.0740740740741
],
[
1052.5925925925926,
298.0740740740741
]
],
"group_id":null,
"shape_type":"polygon",
"flags":{}
},
{
"label":"Footer",
"points":[
[
1033.7407407407406,
1634.5185185185185
],
[
1033.7407407407406,
1646.5185185185185
],
[
1052.7407407407406,
1646.5185185185185
],
[
1052.7407407407406,
1634.5185185185185
]
],
"group_id":null,
"shape_type":"polygon",
"flags":{}
}
],
"imagePath":"val_0031.jpg",
"imageData":null,
"imageHeight":1754,
"imageWidth":1240
}
注意，CDLA的标注的shape_type为polygon，而不是目标检测标注种经常使用的rectangle，所以这导致他有多个points，而Yolo格式需要的是
label bbox中心点横坐标与图像宽度比值 bbox中心点纵坐标与图像高度比值 bbox宽度与图像宽度比值 bbox高度与图像高度比值
所以我们应该先求解多个points的x_min, y_min, x_max, y_max，然后再计算yolo标注数据，
转换脚本如下：
import os
import json
from tqdm import tqdm
def convert_labelme_json_to_yolo_txt(input_json_dir, out_txt_dir, classes):
if not os.path.exists(input_json_dir):
print('json输入文件夹不存在')
return
if not os.path.exists(out_txt_dir):
os.mkdir(out_txt_dir)
# 获取输入文件夹下所有json文件
all_json_files = []
for filename in os.listdir(input_json_dir):
if filename.endswith('.json'):
#file_path = os.path.join(input_json_dir, filename)
#all_json_files.append(file_path)
all_json_files.append(filename)
# 遍历json文件 进行处理
for idx, single_json_filename in tqdm(enumerate(all_json_files)):
single_json_path = os.path.join(input_json_dir, single_json_filename)
print('正在处理第{}个文件{}'.format(idx, single_json_path))
with open(single_json_path, 'r') as single_json_load_file:
json_data = json.load(single_json_load_file)
image_width = json_data['imageWidth']
image_height = json_data['imageHeight']
shapes = json_data['shapes']
out_txt_file_path = os.path.join(out_txt_dir, single_json_filename.replace('json', 'txt'))
if os.path.exists(out_txt_file_path):
try:
os.remove(out_txt_file_path)
print("{}已存在，删除文件".format(out_txt_file_path))
except OSError as e:
print("删除文件{}出错，错误：{}".format(out_txt_file_path, e))
out_txt_file = open(out_txt_file_path, 'w')
for shape_data in shapes:
label = shape_data['label']
labe_index = classes.index(label)
shape_type = shape_data['shape_type']
points = shape_data['points']
x_min, y_min, x_max, y_max = get_polygon_box(points)
# 将标注框按图像大小压缩
x_center = (x_min + x_max) / 2 / image_width
y_center = (y_min + y_max) / 2 / image_height
bbox_w = (x_max - x_min) / image_width
bbox_h = (y_max - y_min) / image_height
bbox = (x_center, y_center, bbox_w, bbox_h)
out_txt_file.writelines(str(labe_index) + " " + " ".join([str(a) for a in bbox]) + '\n')
def get_polygon_box(points):
point_x_list = []
point_y_list = []
for point in points:
point_x = point[0]
point_y = point[1]
point_x_list.append(point_x)
point_y_list.append(point_y)
x_min = min(point_x_list)
x_max = max(point_x_list)
y_min = min(point_y_list)
y_max = max(point_y_list)
return x_min, y_min, x_max, y_max
if __name__ == '__main__':
input_json_path = './CDLA_DATASET/train'
output_txt_path = './CDLA_DATASET_YOLO/train'
classes_list = ['Header', 'Text', 'Reference', 'Figure caption', 'Figure', 'Table caption', 'Table', 'Title', 'Footer', 'Equation']
convert_labelme_json_to_yolo_txt(
input_json_dir=input_json_path,
out_txt_dir=output_txt_path,
classes=classes_list
)
input_json_path = './CDLA_DATASET/train'
output_txt_path = './CDLA_DATASET_YOLO/train'
如果要转换val文件夹，则将main函数下的input_json_path、output_txt_path修改为：
input_json_path = './CDLA_DATASET/val'
output_txt_path = './CDLA_DATASET_YOLO/val'
将原始数据集中的图片按train、val分别拷贝到上述的输出文件夹CDLA_DATASET_YOLO中，组成新的yolo格式的数据集，之后将使用这个数据集进行模型训练。
3 创建yolov8 虚拟环境
搭建yolov8官方文档：https://docs.ultralytics.com/quickstart/#install-ultralytics
首先使用以下命令创建conda虚拟环境，并进入虚拟环境
conda create -n yolov8 python=3.10
conda activate yolov8
然后克隆yolov8仓库，安装yolov8
git clone https://github.com/ultralytics/ultralytics
cd ultralytics
pip install -e .
或者直接使用
pip install ultralytics
安装。
然后再安装pytorch gpu版本，我的命令如下
conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia
4 使用yolov8训练
官方文档：https://docs.ultralytics.com/usage/python/
首先将2.3节中生成数据集拷贝到项目目录CDLA_DATASET_yolo下。
然后我们需要自定义一个数据集配置文件，命名为cdla.yaml，然后将以下内容复制到文件中
train: ./CDLA_DATASET_yolo/train
val: ./CDLA_DATASET_yolo/val
nc: 10 # 分类数
names: ['Header', 'Text', 'Reference', 'Figure caption', 'Figure', 'Table caption', 'Table', 'Title', 'Footer', 'Equation']
然后我们新建一个train_cdla.py文件，在文件中写入以下内容
import sys
import os
from ultralytics import YOLO
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
sys.path.insert(0, os.path.dirname(os.getcwd()))
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
def train_model():
# 加载模型
model = YOLO('./cfg/models/v8/yolov8n.yaml')
# 训练模型
model.train(
name='cdla', # 训练运行的名称，用于在项目文件夹中创建子目录，用于存储训练日志和输出
data='./cdla.yaml', # 数据集配置文件的路径
epochs=300,
batch=32, # 用于训练的批处理大小
optimizer='AdamW',
imgsz=640, # 训练的目标图像大小。在输入模型之前，所有图像都会调整为此尺寸
device=0,
save=True, # 支持保存训练检查点和最终模型权重,对于恢复训练或模型部署很有用
save_period=1, # 保存模型检查点的频率，默认值为-1，表示禁用保存功能
seed=1024, # 设置训练的随机种子，确保在相同配置的运行中结果的可重复性,
lr0=0.0001, # 初始学习率（即 SGD=1E-2 ， Adam=1E-3 ）,
warmup_epochs=10, # 学习率预热的纪元数
plots=True, # 生成并保存训练和验证指标的图以及预测示例，从而提供对模型性能和学习进度的可视化见解
)
metrics = model.val()
print('metric : {}'.format(metrics))
if __name__ == '__main__':
train_model()
执行脚本即可开始训练。
训练相关指标的曲线如下：
5 使用预训练权重推理
新建一个predict.py文件，然后在该文件中写入以下代码
import sys
import os
from ultralytics import YOLO
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
sys.path.insert(0, os.path.dirname(os.getcwd()))
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
def inference_model():
model = YOLO('./runs/detect/cdla/weights/best.pt')
results = model.predict(source='./val_0001.jpg', show=True, save=True)
print(results)
if __name__ == '__main__':
inference_model()
其中
model = YOLO('./runs/detect/cdla/weights/best.pt')
中的路径为预训练权重路径。
运行该脚本，会将结果图片保存在项目根目录/runs/detect/predict文件夹下。
推理结果图片如下：
6 总结
从训练的相关指标和最后的推理结果上看，yolov8在CDLA上训练的版面分析模型还是很有效的，之后可以试一试yolov9看怎么样。
本文作者：StubbornHuang
版权声明：本文为站长原创文章，如果转载请注明原文链接！
原文标题：基于CDLA版面分析数据集使用Yolov8进行文档版面分析实战
原文链接：https://www.stubbornhuang.com/3029/
发布于：2024年04月30日 13:32:05
修改于：2024年04月30日 13:39:51
声明：本站所有文章，如无特殊说明或标注，均为本站原创发布。任何个人或组织，在未征得本站同意时，禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益，可联系我们进行处理。
评论
50