Skip to content
BranchesTags

Repository files navigation

YOLO Training Platform

一个完整的 YOLO 目标检测训练平台，提供 CLI 命令行工具和 Web 界面两种使用方式。支持数据标注管理、离线数据增强、训练预设、超参调优（HPO）、模型对比、模型量化导出等功能。

架构概览

├── training_engine/         # Python 训练引擎（CLI）
│   ├── train.py             #   训练入口
│   ├── tune.py              #   超参调优（遗传算法）
│   ├── predict.py           #   推理入口
│   ├── quantize.py          #   模型量化（FP16/INT8）
│   ├── adapter.py           #   YOLO 模型适配层（ultralytics / yolo26）
│   └── presets.py           #   训练预设（小数据集专用）
├── backend/                 # FastAPI 后端
│   ├── routes/              #   API 路由
│   ├── services/            #   业务逻辑
│   ├── tasks/               #   后台训练任务
│   └── store.py             #   数据存储抽象层
├── frontend/                # React + TypeScript 前端
│   └── src/
│       ├── pages/           #   页面组件
│       ├── components/      #   通用组件
│       └── api/             #   API 对接层
├── scripts/                 # 辅助脚本
├── configs/                 # 数据集配置（YAML）
├── docs/                    # 文档
└── rknn/                    # RKNN 芯片部署工具

环境要求

依赖 版本 说明
Python ≥ 3.10 训练引擎 + 后端
Node.js ≥ 20 前端（Vite + React）
Docker（可选） ≥ 24 Docker 模式一键启动全部服务
PostgreSQL（Docker 模式自动提供） ≥ 16 Docker 模式的数据存储

快速开始

最快体验（一键启动）

# 方式 A：本地脚本（零依赖，SQLite）
bash start-local.sh

# 方式 B：Docker 开发模式（PostgreSQL + 热重载）
docker compose -f docker-compose.dev.yml up -d

启动后访问 http://localhost:3000（Docker 生产模式访问 http://localhost:8000）。

更多启动方式

方式 适用场景 命令
start-local.sh 本地快速体验，零 Docker bash start-local.sh
Docker 开发 全栈开发，热重载 docker compose -f docker-compose.dev.yml up -d
Docker 生产 服务器部署 docker compose up -d
Docker 单容器 云平台部署 docker build -f docker/Dockerfile.prod ...
手动启动 自定义开发环境 分步启动前后端
纯 CLI 只训练，不需要界面 python -m training_engine.train ...

详细的启动说明、环境变量配置、端口速查、常见问题排查，见 启动指南 →。

首次使用 Web 界面

  1. 打开 Web 界面，点击 新建项目
  2. 进入项目 → 上传图片（拖拽或摄像头采集）
  3. 选择图片 → 标注（画框 + 选择类别）
  4. 数据准备好后 → 训练（选择模型和参数，开始训练）
  5. 训练完成 → 在 模型 面板查看指标、导出、推理

CLI 命令行使用

基础训练

# GPU/MPS 训练（Apple Silicon Mac 自动检测）
python -m training_engine.train --model yolov8n.pt --data configs/tennis_ball.yaml --epochs 100

# 指定设备
python -m training_engine.train --model yolov8n.pt --data configs/tennis_ball.yaml --epochs 100 --device mps

# 快速验证（5 个 epoch）
python -m training_engine.train --model yolov8n.pt --data configs/tennis_ball.yaml --epochs 5 --validate-only

小数据集训练

针对 10-200 张图片的场景，提供了专用训练预设和细粒度超参控制：

# 使用预设（快速配置）
python -m training_engine.train \
  --model yolov8n.pt \
  --data configs/your_dataset.yaml \
  --preset aggressive

# 手动指定所有超参
python -m training_engine.train \
  --model yolov8n.pt \
  --data configs/your_dataset.yaml \
  --epochs 200 --imgsz 320 --batch 4 \
  --lr0 0.001 --lrf 0.01 --cos-lr \
  --freeze 10 --close-mosaic 0 \
  --weight-decay 0.001 --warmup-epochs 5

# 两阶段训练（先冻结 backbone，再全网络微调）
# 阶段 1
python -m training_engine.train \
  --model yolov8n.pt --data configs/your_dataset.yaml \
  --epochs 100 --imgsz 320 --batch 4 \
  --lr0 0.001 --cos-lr --freeze 10 --close-mosaic 0

# 阶段 2（使用阶段 1 的 best.pt）
python -m training_engine.train \
  --model runs/small_dataset/stage1/train/weights/best.pt \
  --data configs/your_dataset.yaml \
  --epochs 100 --imgsz 320 --batch 4 \
  --lr0 0.0001 --cos-lr --freeze 0 --close-mosaic 0

# 使用 extra_args 开启高级增强
python -m training_engine.train \
  --model yolov8n.pt --data configs/your_dataset.yaml \
  --extra-args '{"mosaic": 1.0, "mixup": 0.5, "copy_paste": 0.3, "scale": 0.9}'

可用预设

预设 适用场景 epochs imgsz lr0 freeze close_mosaic
conservative 快速基线验证 100 320 0.001 10 10
balanced 50-200 张数据集 150 416 0.001 5 5
aggressive 10-50 张极少数据 200 320 0.0005 10 0（全程 mosaic）
default 标准数据集 100 640 0.01 无 10
high_accuracy 追求最高精度 300 640 0.01 无 15

核心超参说明

参数 作用 小数据集建议值 默认值
--lr0 初始学习率 0.001（降低 10x） 0.01
--lrf 最终学习率因子 0.01 0.01
--cos-lr 余弦退火调度 开启 off
--freeze N 冻结 backbone 前 N 层 10 无
--close-mosaic N 第 N 个 epoch 关闭 mosaic 0（不关闭） 10
--weight-decay 权重衰减（正则化） 0.001 0.0005
--warmup-epochs 学习率 warmup 5 3

超参调优（HPO）

使用遗传算法自动搜索最优超参数：

# 命令行
python -m training_engine.tune \
  --model yolov8n.pt \
  --data configs/tennis_ball.yaml \
  --epochs 30 --iterations 100 \
  --preset aggressive

# 或在 Web 界面中 POST /api/v1/training/tune

推理

# 图片/视频推理
python -m training_engine.predict \
  --model runs/tennis/train/weights/best.pt \
  --source demo.jpg --conf 0.25 --save

# 摄像头实时推理
python scripts/camera_inference.py \
  --model runs/tennis/train/weights/best_int8.onnx \
  --camera 0 --conf 0.25

模型量化

# FP16 量化
python -m training_engine.quantize \
  --model runs/tennis/train/weights/best.pt \
  --format fp16

# INT8 量化（动态）
python -m training_engine.quantize \
  --model runs/tennis/train/weights/best.pt \
  --format int8

# INT8 静态量化（需要校准数据）
python scripts/quantize_onnx.py \
  --model runs/tennis/train/weights/best.pt \
  --calibration-data ./calibration_images/ \
  --method entropy

量化效果预期：

指标 FP32 FP16 INT8
模型大小 1x 0.5x 0.25x
推理速度 1x 1.5-2x 2-4x
mAP 损失 基准 <0.5% <1%

Web 界面功能

访问 http://localhost:3000（开发模式）或 http://localhost:8000（生产模式）后可使用：

项目管理

  • 创建/管理多个项目，每个项目独立的数据集和模型
  • 项目卡片显示缩略图、模型数量、最后编辑时间

数据管理

  • 上传：拖拽上传图片、摄像头采集（Webcam / IP Camera）
  • 标注：内置 YOLO 格式标注工具，支持矩形框和多类别
  • 类别管理：自定义类别名称和颜色
  • 数据增强：离线增强生成新样本，自动调整 bounding box
    • 8 种增强方式：翻转、旋转、亮度对比度、色调饱和度、高斯模糊、噪声、遮挡
    • 3 档预设：基础 / 中等 / 强力
    • 后台异步执行，支持取消和实时状态轮询
    • 增强完成显示错误详情

模型训练

  • 可视化配置：base model、epochs、imgsz、batch、单类别模式
  • 训练预设：一键选择 conservative/balanced/aggressive 等预设
  • 超参调优：遗传算法自动搜索最优参数
  • 实时监控：训练进度条、loss 曲线、mAP 指标
  • 任务管理：查看/取消训练和调优任务

模型管理

  • 模型列表：查看所有已训练模型及其指标
  • 模型对比：多选模型并排对比 mAP50、Precision、Recall 等指标
  • 模型详情：查看训练参数、指标曲线
  • 模型导出：导出 ONNX / FP16 / INT8 格式

推理部署

  • 图片推理：上传图片进行检测
  • 模型部署：导出为 TensorRT / ONNX Runtime 部署配置

数据增强详情

支持 8 种离线增强变换，Bbox 自动跟随几何变换：

增强方式 类别 说明
水平翻转 几何 左右镜像，p=0.5
垂直翻转 几何 上下镜像，适合航拍/显微场景
旋转 ±15° 几何 小幅随机旋转
亮度/对比度 像素 模拟光照变化，limit=0.15
色调/饱和度 像素 随机调整 HSV
高斯模糊 像素 模拟对焦/运动模糊
高斯噪声 像素 模拟传感器噪声
随机遮挡 遮挡 Cutout，4 个随机遮挡块

推荐预设组合：

  • 基础：水平翻转 + 亮度对比度
  • 中等：水平翻转 + 旋转 + 亮度对比度 + 高斯模糊
  • 强力：全部 7 种

项目结构

yolo-training/
├── training_engine/           # Python 训练引擎
│   ├── adapter.py             #   YOLO 模型适配层
│   ├── train.py               #   训练 CLI 入口
│   ├── tune.py                #   HPO 超参调优 CLI
│   ├── predict.py             #   推理 CLI
│   ├── quantize.py            #   模型量化 CLI
│   └── presets.py             #   训练预设定义
├── backend/                   # FastAPI 后端
│   ├── main.py                #   应用入口
│   ├── config.py              #   配置（THUMBNAIL_SIZE 等）
│   ├── models.py              #   SQLAlchemy 模型
│   ├── store.py               #   数据存储抽象层
│   ├── dependencies.py        #   权限/依赖注入
│   ├── routes/                #   API 路由
│   │   ├── datasets.py        #     数据集 + 增强 API
│   │   ├── training.py        #     训练 + HPO API
│   │   ├── models.py          #     模型管理 + 对比 API
│   │   ├── projects.py        #     项目管理 API
│   │   └── ws.py              #     WebSocket（训练进度推送）
│   ├── services/              #   业务逻辑
│   │   ├── augmentation_service.py
│   │   ├── dataset_service.py
│   │   ├── training_service.py
│   │   ├── model_service.py
│   │   └── ...
│   └── tasks/                 #   后台训练任务
├── frontend/                  # React + TypeScript
│   └── src/
│       ├── pages/             #   页面
│       │   ├── Workspace.tsx  #     主工作区
│       │   ├── ImageGrid.tsx  #     图片网格
│       │   ├── TrainingPage.tsx
│       │   ├── ModelPanel.tsx
│       │   └── ...
│       ├── components/        #   组件
│       │   ├── AugmentPanel.tsx
│       │   ├── ModelComparePanel.tsx
│       │   ├── AnnotationTool.tsx
│       │   └── ...
│       └── api/               #   API 对接
├── scripts/                   # 辅助脚本
│   ├── download_dataset.py    #   数据集下载
│   ├── gen_synthetic_dataset.py
│   ├── camera_inference.py    #   摄像头推理
│   └── run_pipeline.py        #   批量训练流水线
├── configs/                   # 数据集 YAML 配置
├── docs/                      # 文档
│   ├── startup-guide.md          #   启动方式指南
│   ├── architecture.md
│   ├── small_dataset_training.md
│   └── project_summary.md
├── rknn/                      # RKNN 芯片部署
├── docker-compose.yml         # 生产部署
└── docker-compose.dev.yml     # 开发环境

常用操作速查

# 训练
python -m training_engine.train --model yolov8n.pt --data configs/tennis_ball.yaml --epochs 100

# 小数据集训练
python -m training_engine.train --model yolov8n.pt --data configs/your_data.yaml --preset aggressive

# 超参调优
python -m training_engine.tune --data configs/tennis_ball.yaml --epochs 30 --iterations 100

# 推理
python -m training_engine.predict --model best.pt --source demo.jpg --save

# 量化
python -m training_engine.quantize --model best.pt --format int8

# 查看训练结果
cat runs/*/train*/results.csv | tail -5

# 启动 Web 界面
bash start-local.sh
# 或: docker compose -f docker-compose.dev.yml up -d

About

No description, website, or topics provided.

Resources

Readme
Activity
Custom properties

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages