YoloTrain - YOLO视觉模型训练平台

YoloTrain是一个可视化管理yolo视觉模型训练的系统，为计算机视觉任务提供了直观的图形界面。该平台集成了数据集管理、模型管理、训练管理和目标检测功能，支持windows、linux、docker等多种部署方式，使用户能够轻松地训练和部署YOLO模型，支持CPU和GPU，使用tensorboard实时查看训练进度，具备数据集自动分割、数据集增强、实时检测等功能。

• 当前版本：目前暂不支持yolov5，支持yolov8以上
• 开发阶段：本项目处于早期开发阶段，不建议用于生产环境
• 开源协议：本项目独立开发的功能完全开源、不限制商用，训练推理功能使用Ultralytics，需遵循AGPL-3.0协议
• Python版本：需Python 3.12，推荐环境：torch2.6.0、CUDA12.6

1. 安装PostgreSQL（用户名postgres，密码postgres）
2. 安装Python 3.12并添加到PATH
3. 克隆项目并运行启动.bat
4. 访问 http://localhost:8000

# CPU版本
docker-compose up -d

# GPU版本（需安装NVIDIA Container Toolkit）
docker-compose -f docker-compose-gpu.yml up -d

如需GPU训练，手动安装torch：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

• 支持多种导入方式（系统数据集、本地上传ZIP、服务器目录）
• AI自动标注功能，支持目标检测、实例分割、OBB旋转框、Pose关键点
• 一键导出为系统数据集或标签文件
• 支持快捷键切换、辅助线、标注框编辑等功能

• 支持ZIP上传和服务器目录导入
• 自动验证数据集结构，识别类别和图片数量
• 数据集比例分割（训练集/验证集/测试集）
• 支持COCO格式JSON转换

• 上传和管理预训练模型及训练后的模型
• 支持YOLOv8/v11等多种模型
• 模型剪枝功能（非结构化和结构化剪枝）

• 创建和管理训练任务
• 支持GPU、CPU、华为昇腾NPU多种设备
• 多GPU/自动GPU训练
• 实时TensorBoard监控训练进度
• 训练日志下载

• 图像和视频目标检测
• 模型对比测试功能
• 可配置置信度阈值、IoU阈值等参数
• 支持目标检测、实例分割、OBB、Pose等任务类型

• 图像预处理（调整大小、去噪、亮度对比度、锐化）
• 图像质量分析
• 数据增强（翻转、旋转、噪声、透视变换）
• 高级增强（CutMix、MixUp、Mosaic、天气模拟）

• 视频提取帧
• 实时检测（支持RTSP、RTMP、HTTP、摄像头）
• 报警区域绘制和多目标报警
• 运动轨迹追踪和预测

YoloTrain/
├── app/                        # 应用程序代码
│   ├── api/                    # API 端点
│   ├── core/                   # 核心配置
│   ├── crud/                   # 数据库 CRUD 操作
│   ├── db/                     # 数据库连接
│   ├── models/                 # 数据库模型
│   ├── schemas/                # Pydantic 模式
│   ├── services/               # 业务逻辑
│   ├── static/                 # 静态文件（CSS/JS/数据集/模型等）
│   └── main.py                 # 应用程序入口点
├── logs/                       # 日志文件（TensorBoard日志）
├── datasets_import/            # 数据集导入目录
├── requirements.txt            # 依赖项
├── run.py                      # 运行脚本
└── Docker相关文件

• 操作系统：Windows 10/11、Linux（Ubuntu 18.04+）或macOS
• Python：3.12或更高版本
• 数据库：PostgreSQL 12或更高版本
• 硬件：
  • CPU：建议4核心以上
  • 内存：建议8GB RAM以上
  • 存储：至少10GB可用空间
  • GPU：支持CUDA的NVIDIA GPU（推荐，但不是必需）

准备YOLO格式的数据集：

dataset/
├── classes.txt           # 类别名称文件
├── train/
│   ├── images/           # 训练图像
│   └── labels/           # 训练标签
├── val/
│   ├── images/           # 验证图像
│   └── labels/           # 验证标签
└── test/                 # 可选
    ├── images/           # 测试图像
    └── labels/           # 测试标签

1. 导入数据集：通过ZIP上传或服务器目录导入
2. 上传模型：上传预训练模型或使用系统预置模型
3. 创建训练任务：选择数据集、模型，配置训练参数
4. 监控训练：通过TensorBoard实时查看训练进度
5. 模型测试：使用训练好的模型进行目标检测

检查是否在项目虚拟环境中，确认python、CUDA、torch版本是否适配。

• 必须有classes.txt或yaml文件
• 图片建议尺寸640×640，宽为8的倍数
• ZIP压缩包内文件必须在根目录

• 确保PostgreSQL服务正在运行
• 检查配置文件config.py和init_db.py中的连接参数

import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA版本: {torch.version.cuda}")
    print(f"GPU数量: {torch.cuda.device_count()}")

感谢所有开发者的帮助与支持

本项目采用AGPL-3.0协议，训练推理功能依赖Ultralytics YOLO。