English | 简体中文
本项目是针对人形机器人椅子检测场景的 YOLO(You Only Look Once)目标检测工具集,提供视频数据预处理、标注辅助和云端训练环境配置等功能。
YOLOchair/ ├── mp4/ # 原始视频文件存储目录 ├── output/ # 处理后的输出文件目录 ├── output_frames_100/ # 视频提取的帧图像(已提取100帧) │ └── frame_000xxx.jpg # 帧图像文件 ├── data/ # 训练数据目录 │ └── chair.v2i.yolo26/ # YOLO格式数据集 │ ├── data.yaml # 数据集配置文件 │ ├── train/ # 训练集 │ │ ├── images/ # 训练图像 │ │ └── labels/ # 训练标签 │ ├── valid/ # 验证集 │ │ ├── images/ # 验证图像 │ │ └── labels/ # 验证标签 │ └── test/ # 测试集 │ ├── images/ # 测试图像 │ └── labels/ # 测试标签 ├── src/ # 源代码目录 │ └── *.py # Python处理脚本 └── README.md # 项目说明文档
# 安装基础依赖
pip install opencv-python numpy
# 安装 YOLOv26 训练依赖
cd src
pip install -r requirements.txt
mp4/ 目录src/ 目录下的处理脚本output_frames_100/ 目录cd src
# 1. 检查数据集
python check_dataset.py
# 2. 开始训练
python train_yolov26.py
# 3. 评估模型
python evaluate_yolov26.py
# 4. 测试推理
python inference_yolov26.py
训练数据已准备完成,位于 data/chair.v2i.yolo26 目录下:
data/chair.v2i.yolo26/ ├── data.yaml # 数据集配置文件 ├── train/ # 训练集 │ ├── images/ # 训练图像(约75张) │ └── labels/ # 训练标注 ├── valid/ # 验证集(约10张) │ ├── images/ # 验证图像 │ └── labels/ # 验证标注 └── test/ # 测试集(约10张) ├── images/ # 测试图像 └── labels/ # 测试标注
数据集配置文件 (data/chair.v2i.yolo26/data.yaml):
train: ../data/chair.v2i.yolo26/train/images
val: ../data/chair.v2i.yolo26/valid/images
test: ../data/chair.v2i.yolo26/test/images
nc: 1 # 类别数量
names: ['chair'] # 类别名称
# 创建虚拟环境(推荐)
python -m venv venv
source venv/bin/activate # Linux/Mac
# 或
venv\Scripts\activate # Windows
# 安装 PyTorch(根据您的CUDA版本选择)
# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# CPU版本(不推荐用于训练)
pip install torch torchvision torchaudio
pip install ultralytics
# 安装其他依赖
pip install opencv-python numpy pillow pyyaml
# 检查 PyTorch 和 CUDA
python -c "import torch; print(f'PyTorch版本: {torch.__version__}'); print(f'CUDA可用: {torch.cuda.is_available()}'); print(f'CUDA版本: {torch.version.cuda if torch.cuda.is_available() else \"N/A\"}')"
python -c "import ultralytics; print(f'Ultralytics版本: {ultralytics.__version__}')"
在开始训练前,先检查数据集是否正确:
# check_dataset.py
import yaml
from pathlib import Path
# 读取数据集配置
data_yaml_path = 'data/chair.v2i.yolo26/data.yaml'
with open(data_yaml_path, 'r') as f:
data_config = yaml.safe_load(f)
# 检查各数据集大小
print("=== 数据集检查 ===")
for split in ['train', 'val', 'test']:
key = split if split in data_config else f'{split}_images' if 'images' in data_config else None
if key:
images_dir = Path(data_config[key])
if images_dir.exists():
num_images = len(list(images_dir.glob('*.jpg'))) + len(list(images_dir.glob('*.png')))
print(f"{split.upper()}: {num_images} 张图像")
else:
print(f"{split.upper()}: 图像目录不存在 - {images_dir}")
print(f"\n类别数量: {data_config['nc']}")
print(f"类别名称: {data_config['names']}")
运行检查脚本:
python check_dataset.py
在项目根目录创建 train.py:
# train.py
from ultralytics import YOLO
import torch
def main():
# 检查CUDA是否可用
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"使用设备: {device}")
# 加载预训练模型
# 可选模型: yolov8n, yolov8s, yolov8m, yolov8l, yolov8x
# n: nano(最快,精度最低)
# s: small(推荐)
# m: medium
# l: large
# x: xlarge(最慢,精度最高)
model = YOLO("yolov8n.pt")
# 开始训练
results = model.train(
data="data/chair.v2i.yolo26/data.yaml", # 数据集配置文件
epochs=100, # 训练轮数
imgsz=640, # 图像尺寸
batch=16, # 批次大小(根据GPU显存调整)
device=device, # 设备
workers=4, # 数据加载线程数
name="chair_detection", # 实验名称
project="runs/train", # 保存目录
exist_ok=True, # 允许覆盖
patience=50, # 早停耐心值(50轮无改善则停止)
save=True, # 保存检查点
plots=True, # 生成训练图表
verbose=True, # 显示详细信息
# 数据增强参数
hsv_h=0.015, # 色调增强
hsv_s=0.7, # 饱和度增强
hsv_v=0.4, # 明度增强
degrees=0.0, # 旋转角度
translate=0.1, # 平移
scale=0.5, # 缩放
shear=0.0, # 剪切
perspective=0.0, # 透视变换
flipud=0.0, # 上下翻转概率
fliplr=0.5, # 左右翻转概率
mosaic=1.0, # 马赛克增强概率
mixup=0.0, # 混合增强概率
# 优化器参数
lr0=0.01, # 初始学习率
lrf=0.01, # 最终学习率因子
momentum=0.937, # SGD动量
weight_decay=0.0005, # 权重衰减
warmup_epochs=3.0, # 预热轮数
warmup_momentum=0.8, # 预热动量
warmup_bias_lr=0.1, # 预热偏置学习率
)
print("\n=== 训练完成 ===")
print(f"训练结果保存在: {results.save_dir}")
print(f"最佳模型: {results.save_dir}/weights/best.pt")
print(f"最新模型: {results.save_dir}/weights/last.pt")
if __name__ == "__main__":
main()
---
### 📊 训练结果评估
#### 步骤 4: 评估模型性能
创建评估脚本 `evaluate.py`:
```python
# evaluate.py
from ultralytics import YOLO
def main():
# 加载最佳模型
model = YOLO("runs/train/chair_detection/weights/best.pt")
# 在验证集上评估
metrics = model.val(
data="data/chair.v2i.yolo26/data.yaml",
imgsz=640,
batch=16,
conf=0.25, # 置信度阈值
iou=0.6, # IOU阈值
plots=True, # 生成评估图表
save_json=True, # 保存JSON格式的结果
)
# 打印主要指标
print("\n=== 评估结果 ===")
print(f"mAP50: {metrics.box.map50:.4f}")
print(f"mAP50-95: {metrics.box.map:.4f}")
print(f"精度 (Precision): {metrics.box.mp:.4f}")
print(f"召回率 (Recall): {metrics.box.mr:.4f}")
if __name__ == "__main__":
main()
运行评估脚本:
python evaluate.py
创建推理脚本 inference.py:
# inference.py
from ultralytics import YOLO
import cv2
import os
def main():
# 加载训练好的模型
model = YOLO("runs/train/chair_detection/weights/best.pt")
# 测试图像路径
test_images = [
"data/chair.v2i.yolo26/test/images/frame_000077_jpg.rf.daa0cf136dedd8aa8519c40ad98cb7ef.jpg",
"data/chair.v2i.yolo26/test/images/frame_000084_jpg.rf.3a433888a191ebea338061b2b0219e98.jpg",
]
# 创建输出目录
output_dir = "runs/detect/predictions"
os.makedirs(output_dir, exist_ok=True)
# 对每张图像进行推理
for img_path in test_images:
if not os.path.exists(img_path):
print(f"警告: 图像不存在 - {img_path}")
continue
# 执行推理
results = model(img_path, conf=0.25, iou=0.45)
# 可视化结果
for i, result in enumerate(results):
# 绘制检测结果
annotated_frame = result.plot()
# 保存结果
img_name = os.path.basename(img_path)
output_path = os.path.join(output_dir, f"pred_{img_name}")
cv2.imwrite(output_path, annotated_frame)
print(f"\n=== {img_name} ===")
for box in result.boxes:
cls_id = int(box.cls[0])
conf = float(box.conf[0])
class_name = result.names[cls_id]
print(f"类别: {class_name}, 置信度: {conf:.2f}")
print(f"\n推理结果已保存至: {output_dir}")
if __name__ == "__main__":
main()
运行推理脚本:
python inference.py
| 参数 | 说明 | 推荐值 | 调整建议 |
|---|---|---|---|
epochs | 训练轮数 | 100 | 数据少时减少,数据多时增加 |
imgsz | 图像尺寸 | 640 | 显存充足时可增大至1280 |
batch | 批次大小 | 16 | 根据GPU显存调整(8/16/32/64) |
lr0 | 初始学习率 | 0.01 | 收敛困难时减小至0.001 |
patience | 早停耐心值 | 50 | 根据验证集改善情况调整 |
场景1: 快速验证训练(适用于快速迭代)
model.train(
epochs=20, # 减少训练轮数
batch=32, # 增大批次
imgsz=640,
patience=10, # 快速早停
save_period=5, # 每5轮保存一次
)
场景2: 高精度训练(适用于正式训练)
model.train(
epochs=300, # 增加训练轮数
batch=8, # 减小批次以使用更大模型
imgsz=1280, # 使用更大图像尺寸
patience=100, # 增加早停耐心值
model="yolov8m.pt", # 使用更大的模型
)
场景3: 小数据集训练(适用于数据量少的情况)
model.train(
epochs=50, # 适中的训练轮数
batch=16,
imgsz=640,
augment=True, # 启用数据增强
mosaic=1.0, # 马赛克增强
mixup=0.15, # 混合增强
copy_paste=0.3, # 复制粘贴增强
)
解决方案:
# 方法1: 减小批次大小
model.train(batch=8) # 或更小,如 4, 2
# 方法2: 减小图像尺寸
model.train(imgsz=480)
# 方法3: 使用更小的模型
model = YOLO("yolov8n.pt") # 使用nano版本
解决方案:
# 降低学习率
model.train(lr0=0.001, lrf=0.01)
# 检查数据集标注质量
# 确保标注框格式正确(YOLO格式: class x_center y_center width height,均为相对坐标)
# 增加训练轮数
model.train(epochs=200, patience=100)
解决方案:
# 增加数据加载线程数
model.train(workers=8)
# 使用更小的模型
model = YOLO("yolov8n.pt")
# 减小图像尺寸
model.train(imgsz=480)
# 启用缓存(首次运行会慢,后续会快)
model.train(cache="ram") # 缓存到内存
# 或
model.train(cache="disk") # 缓存到磁盘
解决方案:
# 确保 data.yaml 中的路径正确
# 检查 data.yaml 内容:
train: ../data/chair.v2i.yolo26/train/images
val: ../data/chair.v2i.yolo26/valid/images
test: ../data/chair.v2i.yolo26/test/images
# 确保路径相对位置正确
# 如果从项目根目录运行,路径应该是相对于项目根目录的
解决方案:
# 使用更大的模型
model = YOLO("yolov8m.pt") # 或 yolov8l.pt, yolov8x.pt
# 增加训练轮数
model.train(epochs=300)
# 调整数据增强
model.train(
mosaic=1.0,
mixup=0.15,
copy_paste=0.3,
degrees=10.0, # 增加旋转
scale=0.9, # 增加缩放
)
# 收集更多训练数据
# 当前训练集约75张图像,建议至少200-500张
# export.py
from ultralytics import YOLO
def main():
# 加载训练好的模型
model = YOLO("runs/train/chair_detection/weights/best.pt")
# 导出为ONNX格式
model.export(format="onnx", opset=12)
# 导出为TensorRT引擎(需要CUDA)
# model.export(format="engine", device=0)
# 导出为TFLite(适用于移动设备)
# model.export(format="tflite")
print("模型导出完成!")
if __name__ == "__main__":
main()
# example_usage.py
from ultralytics import YOLO
def main():
# 加载模型
model = YOLO("runs/train/chair_detection/weights/best.pt")
# 推理
results = model("your_image.jpg")
# 处理结果
for result in results:
boxes = result.boxes
for box in boxes:
cls_id = int(box.cls[0])
conf = float(box.conf[0])
x1, y1, x2, y2 = box.xyxy[0].tolist()
class_name = result.names[cls_id]
print(f"检测到: {class_name} (置信度: {conf:.2f}) 位置: [{x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f}]")
if __name__ == "__main__":
main()
支持将视频文件按固定帧数或时间间隔提取为图像序列:
为 YOLO 模型训练准备标准化的图像数据集:
本项目支持云端训练环境配置,可在以下平台使用:
相关配置脚本请查看 src/ 目录。
本项目仅供学习和研究使用。