算数验证码识别模型训练

---

主题

基于 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

训练步骤

1. 数据集收集
2. 数据集标注
3. 数据集格式转化
4. 数据集训练
5. 模型使用

数据集收集

编写爬虫采集验证码图片，图片集私信领取

数据集标注

每张图片标注 目标对象的位置和对象分类
本例分类为 0-9，+, -, x, =, ? 共15类

1. 命令行启动标注工具labelme

labelme

2. 选择数据集
   点击【打开目录】，选择待标注的图片文件夹（收集的数据集文件夹）
3. 开启自动保存
   点击【文件】- 选中【自动保存】
4. 标注单张图片的步骤如下

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

效果如下