first commit
This commit is contained in:
douboer
332
README.md
332
README.md
@@ -1,102 +1,284 @@
|
||||
# 项目说明
|
||||
# YOLO数字识别系统
|
||||
|
||||
本仓库提供一个基于 TypeScript 的命令行工具，用于识别带有双条干扰线的四位数字验证码。默认情况下，`train/` 目录用于训练，`valid/` 目录用于验证。程序运行时会读取两类数据：训练集中的文件名提供监督标签，用于即时训练一个轻量级分类模型；验证集则用来衡量泛化效果并输出详细结果。
|
||||
使用YOLOv8模型识别图片中的4位阿拉伯数字。经过图片预处理优化，识别准确率显著提升。
|
||||
|
||||
## 环境要求
|
||||
## 🎯 项目亮点
|
||||
|
||||
- Node.js 18 及以上版本（开发环境使用的是 Node.js v22）
|
||||
- 已安装的系统依赖：`sharp` 依赖 libvips，macOS / Linux 上安装该库后方可正常编译
|
||||
- ✅ **高准确率**: 经过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
|
||||
npm install
|
||||
pip install ultralytics opencv-python
|
||||
```
|
||||
|
||||
该命令会安装以下主要依赖：
|
||||
|
||||
- `tesseract.js`：用于对比基准识别效果（未经过干扰线处理时的识别结果）。
|
||||
- `sharp`：完成灰度化、归一化、阈值化、裁剪与缩放等图像预处理。
|
||||
- `ts-node` / `typescript`：支持直接运行 TypeScript 源代码。
|
||||
|
||||
## 快速开始
|
||||
|
||||
默认使用 `train/` 训练，并在 `valid/` 上验证：
|
||||
### 2. 一键运行（完整流程）
|
||||
|
||||
```bash
|
||||
npm run ocr
|
||||
python scripts/run_all.py
|
||||
```
|
||||
|
||||
如需指定数据集，可使用如下方式：
|
||||
这将自动执行：
|
||||
1. 从COCO格式转换为YOLO格式数据集
|
||||
2. 训练YOLOv8模型（100轮）
|
||||
3. 在valid文件夹上进行识别
|
||||
|
||||
### 3. 分步执行
|
||||
|
||||
#### 步骤1：准备数据集
|
||||
|
||||
```bash
|
||||
# 指定训练与验证目录
|
||||
npm run ocr -- ./my-train ./my-valid
|
||||
|
||||
# 仅指定验证目录（训练仍使用默认 train/）
|
||||
npm run ocr -- ./my-valid
|
||||
python scripts/prepare_yolo_dataset.py \
|
||||
--root digit-validation \
|
||||
--out yolo_dataset \
|
||||
--val-ratio 0.2 \
|
||||
--seed 20240305
|
||||
```
|
||||
|
||||
当仅提供一个自定义路径且该路径缺少标签（文件名不含四位数字）时，脚本会把它视作新的验证目录，同时保留默认训练集。命令执行结束后，终端会分别输出训练集与验证集的准确率、逐文件预测，以及 Tesseract.js 的识别结果以供对比。
|
||||
参数说明：
|
||||
- `--root`: COCO数据集根目录
|
||||
- `--out`: YOLO数据集输出目录
|
||||
- `--val-ratio`: 验证集比例（默认0.2）
|
||||
- `--seed`: 随机种子（用于可重复划分）
|
||||
|
||||
## 工具工作流程
|
||||
#### 步骤2：训练模型
|
||||
|
||||
1. **扫描数据集**
|
||||
按目录读取所有扩展名为 `.png`、`.jpg`、`.jpeg`、`.bmp` 的图片文件。文件名中的连续四位数字作为标签，若文件名缺少四位数字，则只生成预测，不计入准确率。
|
||||
|
||||
2. **图像预处理**
|
||||
- 使用 `sharp` 将图片 resize 到固定高度（120px），保持纵横比。
|
||||
- 转换为灰度图并做 normalize。
|
||||
- 应用固定阈值生成二值化掩膜，用于定位数字区域与干扰线。
|
||||
|
||||
3. **分割四个数字**
|
||||
- 统计每列、每行的黑色像素数量，找出真正包含数字墨迹的区域，并裁掉纯白边缘。
|
||||
- 由于图片始终包含四个数字，横向等分为四段，再根据列统计结果向内收缩，确保裁剪框贴近数字并尽量避开干扰线。
|
||||
- 对每个数字区域加入细小的边缘留白，并缩放到 `20×20` 像素，形成 400 维的浮点特征向量（像素值归一化到 0~1，黑色越接近 1）。
|
||||
|
||||
4. **模型训练**
|
||||
- 所有数字特征与文件名标签构成训练集。
|
||||
- 使用自实现的多分类 softmax（逻辑回归）模型，采用随机顺序的批量梯度下降，训练 1000 个 epoch。由于数据量小且特征维度低，训练耗时通常不足 1 秒。
|
||||
|
||||
5. **验证码识别与验证**
|
||||
- 使用训练后的权重分别对训练集和验证集的四位数字进行推断。
|
||||
- 根据标签统计准确率，并输出每张图片的预测详情。
|
||||
- 同时调用一次 Tesseract.js（未做特别预处理），记录其识别文本，方便与自训练模型对比。
|
||||
|
||||
## 目录结构
|
||||
|
||||
```
|
||||
├── package.json npm 配置，包含运行脚本与依赖
|
||||
├── tsconfig.json TypeScript 编译配置
|
||||
├── src/
|
||||
│ └── ocr.ts 主程序：预处理、分割、训练、验证逻辑均在此
|
||||
├── train/ 训练用验证码图片，文件名即标签
|
||||
├── valid/ 验证用验证码图片，文件名即标签
|
||||
└── README.md 项目说明（本文档）
|
||||
```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`: 实验名称
|
||||
|
||||
- **利用文件名作为监督信号**：图片标签直接来自于文件名，不需要额外的标注文件，便于扩充数据集。
|
||||
- **干扰线处理方式**：不是直接删除干扰线，而是通过对列/行墨迹统计裁剪出真正的数字区域。由于干扰线位置基本固定，且覆盖面积较窄，裁剪后得到的数字基本没有残留干扰线。
|
||||
- **模型为何选择 Softmax**：数据量为几十张，模型复杂度越低越稳定。逻辑回归 + 400 维像素特征即可达到 100% 训练准确率，且推理速度极快。
|
||||
- **Tesseract.js 调用**：保留这一步仅为记录传统 OCR 在原始图片上的表现，可作为质量对比或回归基线。
|
||||
训练完成后，最佳模型保存在：`runs/digit_yolo/exp1/weights/best.pt`
|
||||
|
||||
## 常见问题
|
||||
#### 步骤3：识别数字（使用最佳模型）
|
||||
|
||||
| 问题 | 解决方案 |
|
||||
| --- | --- |
|
||||
| 运行时报 `Module not found: sharp` 或安装 `sharp` 失败 | 确认系统已安装 libvips。macOS 可使用 `brew install vips`，Linux 可通过发行版包管理器安装。 |
|
||||
| 输出中 `predicted` 与 `expected` 均为空 | 检查文件名是否包含连续四位数字，脚本只会对这样的文件进行训练和验证。 |
|
||||
| 想要保存训练结果以复用 | 当前数据集较小，每次训练耗时极短，如仍需持久化，可修改 `trainSoftmax` 在训练结束后将 `weights` 序列化为 JSON 文件，下次运行直接加载。 |
|
||||
| 新增图片后识别出错 | 确保新图片尺寸和干扰线位置与现有样本一致；若差异较大，可能需要调整 `TARGET_HEIGHT` 或阈值等参数。 |
|
||||
```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`: 保存可视化结果
|
||||
|
||||
- 如果干扰线形态发生变化、或字体有明显差异，可进一步加入自定义去噪步骤，例如基于曲线拟合的线条擦除或使用形态学操作。
|
||||
- 若数据集显著扩大，可以将 softmax 替换为轻量级的卷积神经网络（可基于 TensorFlow.js），但在当前数据规模下并非必要。
|
||||
- 可以增加命令行参数，用于切换阈值、输出预测概率、导出训练权重等。
|
||||
## 📊 查看结果
|
||||
|
||||
## 反馈
|
||||
### 训练结果
|
||||
|
||||
如需调整识别策略或扩展功能，可在 `src/ocr.ts` 中直接修改对应逻辑。代码整体结构保持模块化：数据加载 → 特征提取 → 训练 → 推理，便于在任一步骤插入新的处理流程。
|
||||
训练指标保存在 `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
|
||||
|
||||
Reference in New Issue
Block a user