主题
基于 YOLO 训练算术验证码识别模型
验证码示例
环境准备
Python版本
Python>=3.8
环境安装
pip install ultralytics
pip install torch==2.3.1
pip install torchvision==0.18.1
pip install onnxruntime
pip install onnx
pip install labelme
pip install labelme2yolo
训练步骤
- 数据集收集
- 数据集标注
- 数据集格式转化
- 数据集训练
- 模型使用
数据集收集
编写爬虫采集验证码图片
数据集标注
每张图片标注 目标对象的位置和对象分类
本例分类为 0-9,+, -, x, =, ? 共15类
- 命令行启动标注工具labelme
labelme
- 选择数据集
点击【打开目录】,选择待标注的图片文件夹(收集的数据集文件夹)
- 开启自动保存
点击【文件】- 选中【自动保存】
- 标注单张图片的步骤如下
gif上传不成功,这里把gif压缩成了zip文件上传
示例.zip
(2.07 MB, 下载次数: 9)
5.重复步骤4,标注完所有的图片
数据集格式转化
将labelme的json格式标注转换为YOLO支持txt格式的标注
在命令行运行命令
labelme2yolo --json_dir <path/to/images_dir> --val_size 0.1 --test_size 0.1
- 其中json_dir指定图片文件夹
- yolo训练会将图片分为 train(训练集)、val(验证集)、test(测试集) 共3个数据集
- val_size 指定val数据集占比 (0.1表示10%)
- test_size 指定test数据集占比 (0.1表示10%)
- 其余为train数据集占比
训练集(Train Set)是模型学习的主要数据来源
验证集(Validation Set)用于在训练过程中评估模型的性能
测试集(Test Set)用于在模型训练完成后进行最终的评估
输出文件结构如下
模型训练
训练
新建文件train.py
from ultralytics import YOLO
# yolo提供的预训练模型
model = YOLO("yolov8s.pt")
results = model.train(data="YOLODataset/dataset.yaml",
# 最大训练批次
epochs=200,
# 使用cpu训练
device='cpu',
batch=16,
# 图片尺寸,设置为32的倍数,根据实际图片集尺寸调整
imgsz=224,
name='<your task name>',
# 连续50个批次没有改善则停止训练
patience=50)
# 训练完成后导出为更通用的onnx格式
model.export(format="onnx", imgsz=[224, 224])
更多参数参考 https://docs.ultralytics.com/modes/train/#train-settings
运行train.py训练
python3 train.py
关于模型精度
训练过程中会自动生成run文件夹,同时终端会有如下图的输出
mAP50 与 mAP50-95 表示模型的精度
官方表示模型精度在0.8以上就可以说是一个不错的模型
训练完成之后,run/.../weights文件夹保存训练的模型
- best.onnx: 训练中效果最好的Epoch对应的模型文件, onnx格式
- best.pt: 训练中效果最好的Epoch保留的模型文件, pt格式
- last.pt 训练中最后一个Epoch对应的模型文件,pt格式
模型使用
from ultralytics import YOLO
# Load a model
model = YOLO("<your model pt file>")
# Run batched inference on a list of images
results = model(["image1.jpg", "image2.jpg"]) # return a list of Results objects
# Process results list
for result in results:
boxes = result.boxes # Boxes object for bounding box outputs
result.show() # display to screen
效果如下