📚 课程目录第2章 YOLOv5训练自定义数据集（口罩检测）

2. YOLOv5训练自定义数据集（口罩检测）

任务目标

本案例将使用LabelImg标注口罩佩戴数据集（佩戴/未佩戴两类），基于YOLOv5训练自定义检测模型，最终实现图片与视频中的口罩佩戴检测。完成本案例后，您将获得一个可直接调用的best.pt权重文件及配套检测脚本。

环境与依赖

Python 3.8+，PyTorch ≥1.8
YOLOv5官方仓库（v6.0或v7.0均可）
LabelImg标注工具（pip install labelImg）
硬件建议：NVIDIA GPU（显存≥4GB），CPU可训练但速度极慢

数据集准备与标注

创建数据集根目录mask_dataset/，内部结构如下：

mask_dataset/
├── images/
│   ├── train/    # 训练图片（建议≥300张）
│   └── val/      # 验证图片（建议≥50张）
├── labels/
│   ├── train/    # 训练标注txt
│   └── val/      # 验证标注txt
└── mask.yaml     # 数据集配置文件

使用LabelImg标注：类别填写mask（佩戴）和no_mask（未佩戴），标注格式选择YOLO，保存后自动生成与图片同名的.txt文件。
编写mask.yaml：

# mask.yaml
train: mask_dataset/images/train
val: mask_dataset/images/val
nc: 2
names: ['mask', 'no_mask']

完整训练与检测代码

以下脚本整合了数据集划分、训练启动与检测调用，保存为train_mask.py：

import os
import random
import shutil
from pathlib import Path

# ========== 1. 数据集自动划分（若已手动划分可跳过） ==========
def split_dataset(img_root, val_ratio=0.15):
    """将img_root下所有图片按比例划分为train/val，并生成对应空label目录"""
    all_imgs = list(Path(img_root).glob('*.jpg')) + list(Path(img_root).glob('*.png'))
    random.shuffle(all_imgs)
    val_num = int(len(all_imgs) * val_ratio)

    for split, imgs in [('val', all_imgs[:val_num]), ('train', all_imgs[val_num:])]:
        img_dir = Path(f'mask_dataset/images/{split}')
        lbl_dir = Path(f'mask_dataset/labels/{split}')
        img_dir.mkdir(parents=True, exist_ok=True)
        lbl_dir.mkdir(parents=True, exist_ok=True)
        for img in imgs:
            shutil.copy(img, img_dir / img.name)
            # 若已有标注则复制，否则留空等待标注
            lbl = img.with_suffix('.txt')
            if lbl.exists():
                shutil.copy(lbl, lbl_dir / lbl.name)

# ========== 2. 训练入口 ==========
def train():
    # 克隆YOLOv5仓库（首次运行）
    if not Path('yolov5').exists():
        os.system('git clone https://github.com/ultralytics/yolov5.git')
    os.chdir('yolov5')
    os.system('pip install -r requirements.txt')

    # 开始训练：输入尺寸640，batch=16，epochs=100
    cmd = (
        'python train.py '
        '--data ../mask_dataset/mask.yaml '
        '--weights yolov5s.pt '
        '--img 640 '
        '--batch 16 '
        '--epochs 100 '
        '--name mask_exp'
    )
    os.system(cmd)

# ========== 3. 检测推理 ==========
def detect(source):
    """source可为图片路径、视频路径或0(摄像头)"""
    os.chdir('yolov5')
    cmd = (
        'python detect.py '
        f'--weights runs/train/mask_exp/weights/best.pt '
        f'--source "{source}" '
        '--conf-thres 0.5 '
        '--save-txt'
    )
    os.system(cmd)
    print(f'检测结果已保存至 yolov5/runs/detect/')

if __name__ == '__main__':
    # 首次运行：将原始图片放入raw_imgs/后执行划分
    # split_dataset('raw_imgs')
    train()
    # 训练完成后测试单张图片
    # detect('path/to/test.jpg')

运行与验证

数据准备：将所有原始图片放入raw_imgs/目录，用LabelImg完成标注后，取消注释split_dataset('raw_imgs')执行一次划分。
开始训练：运行python train_mask.py，训练过程中终端会实时显示mAP、loss等指标。
验证通过标准：训练完成后检查yolov5/runs/train/mask_exp/weights/best.pt是否存在，且训练日志中mAP@0.5达到0.85以上（若数据集质量良好）。
效果测试：执行detect('test.jpg')，查看yolov5/runs/detect/下输出图片，观察是否准确框出佩戴口罩/未佩戴口罩的人脸，类别标签正确。

常见问题

显存不足：降低--batch至8或4，或减小--img至416。
标注类别不匹配：LabelImg中类别名称必须与mask.yaml的names完全一致。
训练不收敛：检查标注框是否过小或过大，确保标注框紧贴目标；增加训练轮次至150。