准备工作
1. clone yolov5 的仓库
git clone https://github.com/ultralytics/yolov5.git
2. 准备数据集
我看网上好多文章都是在讲如何标注，自己制作数据集；此处提供一个从外网下载的口罩数据集，已经做好标注
https://www.alipan.com/s/bpR8veU9k48
数据集来源于墙外分享，不得商用
下载解压后，可以看到obj目录内有一系列jpg和txt, 每张图片对应于一个txt， txt的内容描述的是图片的标注信息。此数据集内的标注文件格式都是yolo的格式。
我们需要把这些文件分割成多个子集：
- 训练集: 用于训练的数据集
- 验证集：---验证-----
- 测试集：---测试-----
每个数据集有image和labels两部分，需要分别存储在不同的目录，下面的脚本用于完成这些工作（请酌情修改）:
import torch
# from IPython.display import Image # for displaying images
import os
import random
import shutil
from sklearn.model_selection import train_test_split
import xml.etree.ElementTree as ET
from xml.dom import minidom
from tqdm import tqdm
from PIL import Image, ImageDraw
import numpy as np
import matplotlib.pyplot as plt
import glob
import shutil
random.seed(108)
def make_data_set(path:str):
# get full path
images = glob.glob(os.path.join(path, "*.jpg"))
annotations = glob.glob(os.path.join(path, "*.txt"))
images.sort()
annotations.sort()
# 测试集0.2, 训练集0.8
train_images, val_images, train_annotations, val_annotations = train_test_split(images, annotations, test_size = 0.2, random_state = 1)
# 测试0.1 验证 0.1
val_images, test_images, val_annotations, test_annotations = train_test_split(val_images, val_annotations, test_size = 0.5, random_state = 1)
images_ = ["images/train", "images/val", "images/test"]
annotations_ = ["labels/train", "labels/val", "labels/test"]
for i in range(len(images_)):
if not os.path.exists(images_[i]):
os.makedirs(images_[i], exist_ok=True)
if not os.path.exists(annotations_[i]):
os.makedirs(annotations_[i], exist_ok=True)
for f in train_images:
shutil.copy(f, images_[0])
for f in val_images:
shutil.copy(f, images_[1])
for f in test_images:
shutil.copy(f, images_[2])
for f in train_annotations:
shutil.copy(f, annotations_[0])
for f in val_annotations:
shutil.copy(f, annotations_[1])
for f in test_annotations:
shutil.copy(f, annotations_[2])
这个脚本从obj目录提取图片和标注文件，按照8:1:1的比例分割训练、验证和测试数据集，并且将其分别存储不同的目录中
准备训练所需yaml文件:
在yolov5/data下创建mask_data.yaml, 内容如下：
# 路径根据实际情况填写
train: ../images/train/
val: ../images/val/
test: ../images/test/
# 2个分类
nc: 2
names: ["masked","no_mask"]
开始训练
进入yolov5目录，输入以下命令
python train.py --img 480 --cfg yolov5s.yaml --batch 32 --epochs 100 --data mask_data.yaml --weights yolov5s.pt --name mask_data --device 0
上面要注意的是--device 选项，对于无显卡的设备，需要改为--device cpu 进行训练，使用cpu训练非常慢
使用云平台训练（推荐）
这里用的是矩池云，该平台为每个新用户提供5元体验券，用作本文的实验性训练来说足够了
注册账号，公众号领算力豆，创建实例这些大家按照平台说明操作就好。
我这里选的算力产品是：Pytorch 2.1.1
实例创建完毕后，将数据集上传至服务器并解压，按照前述步骤clone yolov5代码以及编辑好相关配置，执行上文中的训练命令。对于本文的数据集，一般30分钟内即可搞定
使用训练好的权重数据
训练完成后，终端会输出结果保存在哪个目录，我这里是:
runs/train/mask_data7/weights/best.pt
我们可以使用这个权重数据来进行一些简单的测试
python detect.py --source FF.MP4 --weights best.pt --conf-thres 0.25
此处有两点需注意：
- 如果是在windows 上跑detect, 可能会报出PosixPath问题, 需在detect.py文件头部加入:
import pathlib
pathlib.PosixPath = pathlib.WindowsPath
- 如果是没有NV显卡的机器，则需要使用cpu 推理，在上述命令末尾加上
--device cpu即可，使用cpu推理可能会比较慢
推理完成后，终端会输出结果保存目录
