YOLO V8+Python训练手写数字识别

以下是针对​​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'

---

​​常见问题解决​​

1. ​​CUDA内存不足​​：

   • 降低batch大小（如改为batch=8）。
   • 减小imgsz（如imgsz=320）。

2. ​​MNIST检测效果差​​：

   • 修改标注文件，生成更精确的边界框（如用OpenCV轮廓检测）。
   • 改用分类任务（更适合MNIST）：

     bash

     复制

     yolo task=classify mode=train model=yolov8n-cls.pt data=your_data_path epochs=10 imgsz=28

3. ​​Windows路径问题​​：

   • 确保YAML文件中使用/而非\（如path: ./datasets/digits）。

---

​​最终建议​​

• ​​如果目标是纯分类​​：直接用YOLOv8-Cls模式（更简单且适合MNIST）。
• ​​如果需要检测位置​​：换用带真实标注的数据集（如Kaggle上的手写数字检测数据集）。