数据集和代码的准备
yolov5源码下载
数据集的下载
直接下载别人制作好的数据集
然后将其放入与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']
制作个人数据集
我们也可以制作自己的数据集进行训练
数据集标注
首先要安装 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)
训练阶段
模型训练
在运行代码之前,需要搭建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