文章目录

• 数据集和代码的准备
• • yolov5源码下载
  • 数据集的下载
  • 制作个人数据集
  • • 数据集标注
    • json 转 txt
• 训练阶段
• • 模型训练
  • 可视化
• 测试阶段

数据集和代码的准备

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/imagesnc: 2
names: ['mask', 'no-mask']

制作个人数据集

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

数据集标注

首先要安装 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 npdir_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 += strxyftxt.writelines(label + "\n")for cnt, json_name in enumerate(list_json):print('cnt=%d,name=%s' % (cnt, json_name))path_json = dir_json + json_namepath_txt = dir_txt + json_name.replace('.json', '.txt')json2txt(path_json, path_txt)

训练阶段

模型训练

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

然后在pycharm 中切换虚拟环境

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

--data ../MaskDataSet/data.yaml
--cfg models/yolov5s.yaml
--weights weights/yolov5s.pt
--batch-size 16

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