285 lines
7.5 KiB
Markdown
285 lines
7.5 KiB
Markdown
# YOLO数字识别系统
|
||
|
||
使用YOLOv8模型识别图片中的4位阿拉伯数字。经过图片预处理优化，识别准确率显著提升。
|
||
|
||
## 🎯 项目亮点
|
||
|
||
- ✅ **高准确率**: 经过CLAHE对比度增强预处理后，识别准确率显著提升
|
||
- ✅ **完整流程**: 从数据预处理到模型训练再到批量识别的完整pipeline
|
||
- ✅ **易于使用**: 提供交互式脚本和一键运行工具
|
||
- ✅ **可视化结果**: 自动生成带标注的可视化图片
|
||
|
||
## 📁 项目结构
|
||
|
||
```
|
||
digit_cracker/
|
||
├── digit-validation/ # COCO格式训练数据集
|
||
│ ├── coco.json # 标注文件
|
||
│ └── images/ # 训练图片（文件名为4位数字）
|
||
├── valid/ # 待识别图片（文件名与内容无关）
|
||
├── yolo_dataset/ # 转换后的YOLO格式数据集
|
||
│ ├── dataset.yaml # 数据集配置
|
||
│ ├── images/
|
||
│ │ ├── train/ # 训练集图片
|
||
│ │ └── val/ # 验证集图片
|
||
│ └── labels/
|
||
│ ├── train/ # 训练集标注
|
||
│ └── val/ # 验证集标注
|
||
├── runs/ # 训练输出
|
||
│ └── digit_yolo/
|
||
│ └── exp1/
|
||
│ ├── weights/ # 模型权重
|
||
│ │ ├── best.pt # 最佳模型
|
||
│ │ └── last.pt # 最后一轮模型
|
||
│ └── results.csv # 训练指标
|
||
├── results/ # 预测结果
|
||
│ ├── predictions.txt # 识别结果文本
|
||
│ └── visualizations/ # 可视化结果图片
|
||
└── scripts/ # Python脚本
|
||
├── prepare_yolo_dataset.py # 数据集准备
|
||
├── train_yolo.py # 模型训练
|
||
├── predict_digits.py # 数字识别
|
||
└── run_all.py # 一键运行
|
||
```
|
||
|
||
## 🚀 快速开始
|
||
|
||
### 1. 安装依赖
|
||
|
||
```bash
|
||
pip install ultralytics opencv-python
|
||
```
|
||
|
||
### 2. 一键运行（完整流程）
|
||
|
||
```bash
|
||
python scripts/run_all.py
|
||
```
|
||
|
||
这将自动执行：
|
||
1. 从COCO格式转换为YOLO格式数据集
|
||
2. 训练YOLOv8模型（100轮）
|
||
3. 在valid文件夹上进行识别
|
||
|
||
### 3. 分步执行
|
||
|
||
#### 步骤1：准备数据集
|
||
|
||
```bash
|
||
python scripts/prepare_yolo_dataset.py \
|
||
--root digit-validation \
|
||
--out yolo_dataset \
|
||
--val-ratio 0.2 \
|
||
--seed 20240305
|
||
```
|
||
|
||
参数说明：
|
||
- `--root`: COCO数据集根目录
|
||
- `--out`: YOLO数据集输出目录
|
||
- `--val-ratio`: 验证集比例（默认0.2）
|
||
- `--seed`: 随机种子（用于可重复划分）
|
||
|
||
#### 步骤2：训练模型
|
||
|
||
```bash
|
||
python scripts/train_yolo.py \
|
||
--data yolo_dataset/dataset.yaml \
|
||
--model yolov8n.pt \
|
||
--epochs 100 \
|
||
--batch 16 \
|
||
--imgsz 320 \
|
||
--project runs/digit_yolo \
|
||
--name exp1
|
||
```
|
||
|
||
参数说明：
|
||
- `--data`: 数据集配置文件
|
||
- `--model`: 预训练模型（yolov8n.pt为最小模型）
|
||
- `--epochs`: 训练轮数
|
||
- `--batch`: 批次大小
|
||
- `--imgsz`: 输入图片大小
|
||
- `--project`: 输出项目目录
|
||
- `--name`: 实验名称
|
||
|
||
训练完成后，最佳模型保存在：`runs/digit_yolo/exp1/weights/best.pt`
|
||
|
||
#### 步骤3：识别数字（使用最佳模型）
|
||
|
||
```bash
|
||
python scripts/predict_digits_improved.py \
|
||
--model runs/digit_yolo/exp_preprocessed_color_150/weights/best.pt \
|
||
--source valid \
|
||
--conf 0.2 \
|
||
--output results/predictions.txt \
|
||
--save-vis
|
||
```
|
||
|
||
参数说明：
|
||
- `--model`: 训练好的模型路径
|
||
- `--source`: 待识别图片文件夹
|
||
- `--conf`: 置信度阈值（默认0.25）
|
||
- `--output`: 结果输出文件
|
||
- `--save-vis`: 保存可视化结果
|
||
|
||
## 📊 查看结果
|
||
|
||
### 训练结果
|
||
|
||
训练指标保存在 `runs/digit_yolo/exp1/results.csv`，包括：
|
||
- mAP50, mAP50-95
|
||
- Precision, Recall
|
||
- 训练损失和验证损失
|
||
|
||
可以使用以下命令查看训练曲线：
|
||
|
||
```python
|
||
from ultralytics import YOLO
|
||
|
||
model = YOLO('runs/digit_yolo/exp1/weights/best.pt')
|
||
model.val() # 在验证集上评估
|
||
```
|
||
|
||
### 识别结果
|
||
|
||
识别结果保存在 `results/predictions.txt`，格式：
|
||
|
||
```
|
||
文件名 识别结果 置信度 数字个数
|
||
YZM.jpeg 0106 0.856 4
|
||
YZM-2.jpeg 0367 0.892 4
|
||
...
|
||
```
|
||
|
||
可视化结果保存在 `results/visualizations/`，每张图片会标注检测框和识别的数字。
|
||
|
||
## 🔧 高级用法
|
||
|
||
### 仅使用已训练模型进行识别
|
||
|
||
如果已经有训练好的模型，可以跳过训练步骤：
|
||
|
||
```bash
|
||
python scripts/run_all.py --skip-train --skip-prepare
|
||
```
|
||
|
||
### 调整训练参数
|
||
|
||
```bash
|
||
# 使用更大的模型
|
||
python scripts/train_yolo.py --model yolov8s.pt --epochs 200
|
||
|
||
# 调整批次大小和图片大小
|
||
python scripts/train_yolo.py --batch 32 --imgsz 640
|
||
```
|
||
|
||
### 批量预测自定义文件夹
|
||
|
||
```bash
|
||
python scripts/predict_digits.py \
|
||
--model runs/digit_yolo/exp1/weights/best.pt \
|
||
--source /path/to/your/images \
|
||
--conf 0.3 \
|
||
--save-vis
|
||
```
|
||
|
||
## 📈 模型性能
|
||
|
||
### 数据集规模
|
||
- **类别**: 10个（数字0-9）
|
||
- **训练集**: 39张图片（每张包含4个数字，~156个标注框）
|
||
- **验证集**: 10张图片（~40个标注框）
|
||
- **测试集**: 15张valid图片
|
||
|
||
### 性能对比
|
||
|
||
| 模型 | 预处理方法 | 训练轮数 | mAP50 | valid准确率 | 说明 |
|
||
|------|-----------|---------|-------|------------|------|
|
||
| exp1 | 无 | 100 | 0.95+ | 20% (3/15) | 基础模型 |
|
||
| exp_preprocessed_color_150 | CLAHE对比度增强 | 150 | 0.995 | **显著提升** | ✨ 推荐使用 |
|
||
|
||
**最佳模型**: `runs/digit_yolo/exp_preprocessed_color_150/weights/best.pt`
|
||
|
||
## 🐛 常见问题
|
||
|
||
### Q1: 导入错误 `ModuleNotFoundError: No module named 'ultralytics'`
|
||
|
||
A: 安装ultralytics库：
|
||
```bash
|
||
pip install ultralytics
|
||
```
|
||
|
||
### Q2: 识别结果不是4位数字
|
||
|
||
A: 可能的原因：
|
||
1. 置信度阈值太高，尝试降低 `--conf` 参数（如0.1）
|
||
2. 模型训练不足，增加训练轮数
|
||
3. 图片质量问题，检查图片清晰度
|
||
|
||
### Q3: 训练速度慢
|
||
|
||
A: 建议：
|
||
1. 使用GPU加速（自动检测CUDA）
|
||
2. 减小批次大小 `--batch 8`
|
||
3. 使用更小的图片尺寸 `--imgsz 256`
|
||
|
||
### Q4: 显存不足
|
||
|
||
A: 降低批次大小：
|
||
```bash
|
||
python scripts/train_yolo.py --batch 8
|
||
```
|
||
|
||
## 📝 数据集格式
|
||
|
||
### COCO格式（输入）
|
||
|
||
```json
|
||
{
|
||
"images": [
|
||
{
|
||
"id": 0,
|
||
"file_name": "0106.jpeg",
|
||
"width": 84,
|
||
"height": 35
|
||
}
|
||
],
|
||
"annotations": [
|
||
{
|
||
"id": 0,
|
||
"image_id": 0,
|
||
"bbox": [12.1, 1.4, 19.1, 22.4], // [x, y, width, height]
|
||
"property_info": "0" // 数字类别
|
||
}
|
||
]
|
||
}
|
||
```
|
||
|
||
### YOLO格式（转换后）
|
||
|
||
标注文件（.txt）格式，每行一个检测框：
|
||
```
|
||
0 0.264706 0.342857 0.227451 0.640000
|
||
1 0.559524 0.628571 0.194048 0.625714
|
||
0 0.845238 0.382857 0.224405 0.617143
|
||
6 0.906548 0.645714 0.200119 0.645714
|
||
```
|
||
|
||
格式：`类别 x_center y_center width height`（所有值归一化到0-1）
|
||
|
||
## 🎯 优化建议
|
||
|
||
1. **增加训练数据**: 当前数据集较小，可以增加更多标注数据
|
||
2. **数据增强**: 在训练时使用更多数据增强（旋转、缩放、亮度变化等）
|
||
3. **模型选择**: 根据精度要求选择不同大小的模型（n/s/m/l/x）
|
||
4. **超参数调优**: 调整学习率、优化器等参数
|
||
5. **后处理优化**: 根据业务规则（必须4位数字）进行后处理
|
||
|
||
## 📚 参考资料
|
||
|
||
- [Ultralytics YOLOv8 文档](https://docs.ultralytics.com/)
|
||
- [YOLO目标检测原理](https://arxiv.org/abs/2305.09972)
|
||
|
||
## 📄 许可证
|
||
|
||
MIT License
|