YOLO V8+Python训练手写数字识别
本文详细介绍了在Windows11系统下使用Python和YOLOv8进行目标检测的完整流程。主要内容包括:1) MNIST数据集下载并转换为YOLO格式;2) Python环境配置与依赖库安装;3) 数据集目录结构组织和YAML配置文件创建;4) YOLOv8模型训练、验证和测试的详细步骤。文中还提供了常见问题解决方案,如CUDA内存不足处理、检测效果优化建议等。特别指出对于MNIST这样的分类
·
以下是针对Windows 11 + Python环境的详细步骤说明,从数据集整理到模型训练,全部适配YOLOv8流程。
1. 数据集整理(MNIST → YOLO格式)
1.1 下载MNIST数据集
MNIST数据集可通过Python直接下载(无需手动下载):
python
复制
from torchvision.datasets import MNIST
import os
# 自动下载MNIST数据集(图片和标签)
train_data = MNIST(root='./data', train=True, download=True)
test_data = MNIST(root='./data', train=False, download=True)
- 数据会保存在
./data/MNIST/raw/
目录下,包含压缩文件(无需解压,代码自动处理)。
1.2 转为YOLO格式
MNIST默认是分类数据,需转为检测格式(假设数字居中,占图片50%区域):
python
复制
from PIL import Image
import numpy as np
# 创建目录
os.makedirs('./datasets/digits/train/images', exist_ok=True)
os.makedirs('./datasets/digits/train/labels', exist_ok=True)
os.makedirs('./datasets/digits/val/images', exist_ok=True)
os.makedirs('./datasets/digits/val/labels', exist_ok=True)
# 转换训练集
for i, (img, label) in enumerate(train_data):
# 保存图片
img.save(f'./datasets/digits/train/images/{i}.jpg')
# 生成YOLO标签文件(格式:class x_center y_center width height)
with open(f'./datasets/digits/train/labels/{i}.txt', 'w') as f:
f.write(f"{label} 0.5 0.5 0.8 0.8\n") # 假设数字居中,宽高占80%
# 转换验证集(用测试集代替)
for i, (img, label) in enumerate(test_data):
img.save(f'./datasets/digits/val/images/{i}.jpg')
with open(f'./datasets/digits/val/labels/{i}.txt', 'w') as f:
f.write(f"{label} 0.5 0.5 0.8 0.8\n")
关键说明:
- YOLO格式要求坐标归一化(0~1),这里假设数字始终居中并占图片80%区域。
- 如果需要真实边界框,需用OpenCV生成轮廓或换用已标注的数据集(如Kaggle上的
handwritten-digits-and-letters
)。
2. 系统环境安装与配置
2.1 安装Python
- 从Python官网下载Python 3.8+(勾选
Add to PATH
)。 - 安装后验证:
bash
复制
python --version pip --version
2.2 安装依赖库
bash
复制
# 创建虚拟环境(推荐)
python -m venv yolov8_env
yolov8_env\Scripts\activate # 激活环境
# 安装核心库
pip install ultralytics torch torchvision opencv-python
注意:
- 如果GPU可用,安装CUDA版的PyTorch(需先安装NVIDIA驱动):
bash
复制
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
3. 目录结构
最终目录结构如下:
复制
datasets/
└── digits/
├── train/
│ ├── images/ # 存放MNIST训练图片(如0.jpg, 1.jpg...)
│ └── labels/ # 存放对应的.txt标签文件
└── val/
├── images/ # 验证集图片
└── labels/
digits.yaml # 数据集配置文件
3.1 创建YAML配置文件
新建digits.yaml
:
yaml
复制
path: ./datasets/digits
train: train/images
val: val/images
names:
0: 0
1: 1
2: 2
3: 3
4: 4
5: 5
6: 6
7: 7
8: 8
9: 9
4. 模型训练
4.1 启动训练
在终端运行:
bash
复制
yolo task=detect mode=train model=yolov8n.pt data=digits.yaml epochs=20 imgsz=640 batch=16
参数解释:
model=yolov8n.pt
:使用YOLOv8 Nano模型(轻量级)。imgsz=640
:输入图像尺寸(MNIST原图28x28会被上采样)。batch=16
:批次大小(根据GPU内存调整,若报错则降低为8或4)。epochs=20
:迭代次数(MNIST数据简单,10~20足够)。
4.2 训练结果
- 训练日志和模型权重保存在
runs/detect/train/
目录。 - 查看验证指标(如mAP):
bash
复制
yolo val model=runs/detect/train/weights/best.pt data=digits.yaml
4.3 测试模型
bash
复制
# 测试单张图片
yolo predict model=runs/detect/train/weights/best.pt source='test.jpg'
常见问题解决
-
CUDA内存不足:
- 降低
batch
大小(如改为batch=8
)。 - 减小
imgsz
(如imgsz=320
)。
- 降低
-
MNIST检测效果差:
- 修改标注文件,生成更精确的边界框(如用OpenCV轮廓检测)。
- 改用分类任务(更适合MNIST):
bash
复制
yolo task=classify mode=train model=yolov8n-cls.pt data=your_data_path epochs=10 imgsz=28
-
Windows路径问题:
- 确保YAML文件中使用
/
而非\
(如path: ./datasets/digits
)。
- 确保YAML文件中使用
最终建议
- 如果目标是纯分类:直接用YOLOv8-Cls模式(更简单且适合MNIST)。
- 如果需要检测位置:换用带真实标注的数据集(如Kaggle上的手写数字检测数据集)。

GitCode 天启AI是一款由 GitCode 团队打造的智能助手,基于先进的LLM(大语言模型)与多智能体 Agent 技术构建,致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话,还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力,真正做到“一句话,让 Al帮你完成复杂任务”。
更多推荐
所有评论(0)