YOLOv11 交通违规检测数据集训练方案：头盔佩戴与闯红灯识别

推荐目录结构

traffic_violation_dataset/
├── images/
│   ├── train/ # 4742 张
│   └── val/ # 390 张
├── labels/
│   ├── train/ # 对应的 .txt 标签文件
│   └── val/
├── traffic.yaml # YOLO 配置文件
└── train.py # 训练主程序

YOLO 配置文件：traffic.yaml

# traffic.yaml
train: ./traffic_violation_dataset/images/train
val: ./traffic_violation_dataset/images/val
nc: 7
names: ['Crossing_Violation', 'Crosswalk_Violation', 'Helmet_Violation', 'Normal', 'Passenger_Violation', 'Pedestrian_Violation', 'Trafficlight_Violation']

✅ 将此文件保存在项目根目录下。

YOLOv11 完整训练代码：train.py

# train.py
from ultralytics import YOLO
import os
import time

# ======================= 参数配置 =======================
MODEL_NAME = "yolov11s.pt"  # 推荐使用 yolov11n/s/m；小显存选 n，精度高选 m/l
CONFIG_FILE = "traffic.yaml"  # 数据集配置路径
EPOCHS = 100
BATCH_SIZE = 16  # 若显存不足可设为 8 或 4
IMG_SIZE = 640  # 输入尺寸（可调整为 1280 提升小目标检测能力）
PROJECT_NAME = "traffic_violation_detection"
EXPERIMENT_NAME = "exp_traffic_yolo11"
DEVICE = 0 if os.environ.get("CUDA_AVAILABLE") else "cpu"  # 自动选择 GPU/CPU

# 检查配置文件是否存在
if not os.path.exists(CONFIG_FILE):
    raise FileNotFoundError(f"[ERROR] 找不到配置文件：{CONFIG_FILE}")

# ======================= 加载模型 =======================
print("🚀 加载 YOLOv11 模型...")
model = YOLO(MODEL_NAME)  # 自动下载预训练权重

# ======================= 开始训练 =======================
print("🔥 开始训练交通违规检测模型...")
start_time = time.time()
results = model.train(
    data=CONFIG_FILE,
    epochs=EPOCHS,
    batch=BATCH_SIZE,
    imgsz=IMG_SIZE,
    device=DEVICE,
    project=PROJECT_NAME,
    name=EXPERIMENT_NAME,
    exist_ok=True,
    patience=30,  # 早停机制：验证损失连续 30 轮不降则停止
    save=True,
    save_period=10,  # 每 10 轮保存一次检查点
    cache=False,  # 设置为 True 可加快加载速度（需足够内存）
    workers=4,
    # 优化器设置
    optimizer='AdamW',
    lr0=0.001,  # 初始学习率
    lrf=0.01,  # 最终学习率 = lr0 * lrf
    momentum=0.937,
    weight_decay=0.0005,
    # 数据增强策略（适合交通场景）
    hsv_h=0.015,  # 色调扰动
    hsv_s=0.7,
    hsv_v=0.4,
    degrees=10.0,  # 轻微旋转模拟摄像头角度变化
    translate=0.2,
    scale=0.5,
    shear=2.0,
    flipud=0.0,  # 上下翻转不合理（天空/地面不对称）
    fliplr=0.5,  # 左右翻转增强
    mosaic=1.0,  # Mosaic 增强提升小目标和遮挡鲁棒性
    mixup=0.1,  # MixUp 缓解过拟合
    copy_paste=0.1  # 对 Helmet_Violation 等稀疏类有帮助
)

# ======================= 输出训练耗时 =======================
training_time = (time.time() - start_time) / 60
print(f"✅ 训练完成！总耗时：{training_time:.2f} 分钟")
print(f"📍 模型保存路径：{results.save_dir}")

# ======================= 验证模型性能 =======================
print("🔍 开始验证...")
metrics = model.val()
print(f"[email protected]: {metrics.box.map50:.4f}")
print(f"[email protected]:0.95: {metrics.box.map:.4f}")
print(f"Precision: {metrics.box.p:.4f}")
print(f"Recall: {metrics.box.r:.4f}")
# 显示各类别 AP @0.5
print("\n📊 各类别 [email protected]:")
for i, cls_name in enumerate(model.names):
    print(f" {cls_name}: {metrics.box.maps[i]:.4f}")

# ======================= 导出 ONNX 模型（便于部署） =======================
print("📦 导出 ONNX 模型...")
onnx_path = model.export(format="onnx", dynamic=True, simplify=True)
print(f"ONNX 模型已保存至：{onnx_path}")

推理示例代码：inference.py

# inference.py
from ultralytics import YOLO
import cv2

# 加载训练好的模型
model = YOLO("traffic_violation_detection/exp_traffic_yolo11/weights/best.pt")

# 单张图片推理
image_path = "test.jpg"
results = model(image_path, conf=0.4, iou=0.5, imgsz=640)

# 可视化结果
for r in results:
    im_array = r.plot()  # 绘制检测框和标签
    im = cv2.cvtColor(im_array, cv2.COLOR_RGB2BGR)
    cv2.imshow("Traffic Violation Detection", im)
    cv2.waitKey(0)
    cv2.imwrite("result_detected.jpg", im)

# 打印检测详情
for r in results:
    boxes = r.boxes
    for box in boxes:
        cls_id = int(box.cls[0])
        conf = float(box.conf[0])
        class_name = model.names[cls_id]
        print(f"检测到违规行为：{class_name} (置信度：{conf:.2f})")

应用场景建议

注意事项与优化建议