YOLOv5实操——检测是否戴口罩

数据集和代码的准备

yolov5源码下载

yolov5官方下载地址

数据集的下载

直接下载别人制作好的数据集

Mask Wearing Dataset 下载

然后将其放入与yolov5-master同级的目录下。

打开MaskDataSet ，可以看到其目录结构如下：

images文件夹下的图片

lables文件夹下对应的同名txt文件。

data.yaml ： train 和val 分别制定了训练集和测试集文件夹相对于yolov5-master的路径。nc指定了标签的个数一共两个。names指定了标签的名称分别为mask 和 no-mask

train: ../MaskDataSet/train/images   
val: ../MaskDataSet/valid/images

nc: 2
names: ['mask', 'no-mask']
1
2
3
4
5

制作个人数据集

我们也可以制作自己的数据集进行训练

数据集标注

首先要安装 lableme，在当前的虚拟环境下键入pip install labelme 。安装完成后，再次输入labelme，即可进入labelme的图形化界面。

【file】——>【Change Output Dir】, 选择保存标签的新文件夹（labelme生成的标签是json文件的格式，后续需要转化成txt文件才能被yolov5使用）
【file】——>【Save Automatically】自动保存

快捷键Clt+R ,创建长方形分别框住所有的检测目标。并填写标签（如此处我填写的是mask,代表戴口罩），选择ok 。然后点击【next image】，对其他的图片进行标注。

json 转 txt

标注完成后会生成对应的json文件, 我们需要将其转换为yolov5中需要的txt文件。
转换的脚本如下：

import os
import json
import numpy as np

dir_json = 'E:\\project\\yolo\\data\\json\\'  # json存储的文件目录
dir_txt = 'E:\\project\\yolo\\data\\labels\\'  # txt存储目录
if not os.path.exists(dir_txt):
    os.makedirs(dir_txt)
list_json = os.listdir(dir_json)


def json2txt(path_json, path_txt):  # 可修改生成格式
    with open(path_json, 'r') as path_json:
        jsonx = json.load(path_json)
        with open(path_txt, 'w+') as ftxt:
            for shape in jsonx['shapes']:
                label = str(shape['label']) + ' '
                xy = np.array(shape['points'])
                strxy = ''
                for m, n in xy:
                    m = int(m)
                    n = int(n)
                    strxy += str(m) + ' ' + str(n) + ' '

                label += strxy
                ftxt.writelines(label + "\n")


for cnt, json_name in enumerate(list_json):
    print('cnt=%d,name=%s' % (cnt, json_name))
    path_json = dir_json + json_name
    path_txt = dir_txt + json_name.replace('.json', '.txt')
    json2txt(path_json, path_txt)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

训练阶段

模型训练

在运行代码之前，需要搭建pytorch的虚拟环境。可以参见教程Windows系统用Anaconda安装pytorch的gpu版本

然后在pycharm 中切换虚拟环境

打开train.py文件，配置相关参数

--data ../MaskDataSet/data.yaml
--cfg models/yolov5s.yaml
--weights weights/yolov5s.pt
--batch-size 16
1
2
3
4

debug记录
【1】No module named ‘cv2’
解决方法：pip install opencv-python
【2】RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 4.00 GiB total capacity; 2.56 GiB already allocated; 0 bytes free; 2.57 GiB reserved in total by PyTorch)
解决方法：在参数设置中，将batchsize 调小一点
【3】TypeError: can’t convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
解决方法：点击进入Traceback 中最后一个蓝色链接（tensor.py文件）将self.numpy修改为self.cpu().numpy()

可视化

我们可以使用tensorboard对训练结果进行可视化
首先安装tensorboard pip install tensorboard
然后在pycharm中打开终端 cd yolov5-master, tensorboard --logdir=./runs

在浏览器中查看http://localhost:6006/

测试阶段

添加detect.py的配置文件

--source
./inference/images/
--weights
./weights/best.pt
--conf
0.4
1
2
3
4
5
6