YOLOv4实战尝鲜 — 教你从零开始训练自己的数据集（安全头盔佩戴识别检测）人工智能梦坠凡尘-

本文代码基于：https://github.com/ultralytics/yolov3
YOLOv4理论部分我的另一篇博客请参考：YOLOv4真的来了！！论文翻译

数据准备

首先介绍数据集，来源于AI研习设的一个比赛，见链接：https://god.yanxishe.com/32。
之前我已经用该数据集训练过YOLOv3，感觉效果不是非常好，这次直接把之前训练YOLOv3的数据集放在YOLOv4_path/data/目录下即可。

不过，，这样看的可能会一脸懵逼，还是介绍下数据集的准备过程吧。。

划分数据集

首先clone git 到本地，解压，得到YOLOv4的文件目录如图：

然后将下载下来的数据集放到data/下，文件夹命名为原始数据集。data/原始数据集/下如图所示，分别有train文件夹（训练数据集），labels文件夹（标签集），test文件夹（测试数据，用于测试提交。）

先划分数据集，将train数据集划分为train/valid数据集，运行 1.split_data.py 如下代码：

import os import glob import json import shutil import numpy as np import xml.etree.ElementTree as ET  img_train = './images/train/' img_val = './images/valid/'  label_train = './labels/train/' label_val = './labels/valid/'  allimgs = glob.glob('原始数据集/train/' + "/*.jpg") allimgs = np.sort(allimgs) np.random.seed(100) np.random.shuffle(allimgs)  train_ratio = 0.9 train_num = int(len(allimgs) * train_ratio) # 得到训练和验证数据集列表 img_list_train = allimgs[:train_num] img_list_val = allimgs[train_num:] # 创建文件夹 if os.path.exists(img_train):     shutil.rmtree(img_train)     os.mkdir(img_train) else:     os.mkdir(img_train) if os.path.exists(img_val):     shutil.rmtree(img_val)     os.mkdir(img_val) else:     os.mkdir(img_val) if os.path.exists(label_train):     shutil.rmtree(label_train)     os.mkdir(label_train) else:     os.mkdir(label_train) if os.path.exists(label_val):     shutil.rmtree(label_val)     os.mkdir(label_val) else:     os.mkdir(label_val) # 移动val数据到指定位置 for i in img_list_val:     img_id = i.split('.')[0].split('/')[2] print(img_id) # jpg     shutil.copy(i, img_val + img_id + '.jpg') # xml     shutil.copy('原始数据集/label/' + 'new_' + img_id + '.xml', label_val + img_id + '.xml') # 移动train数据到指定位置 for i in img_list_train:     img_id = i.split('.')[0].split('/')[2] print(img_id) # jpg     shutil.copy(i, img_train + img_id + '.jpg') # xml     shutil.copy('原始数据集/label/' + 'new_' + img_id + '.xml', label_train + img_id + '.xml') 

至此，得到的后面用于训练验证的数据集images和对应labels文件夹。见如下文件夹

xml2yolo

yolo的标签格式是txt格式，所以我们还需将xml标签转为txt格式。先新建txt_train、txt_valid文件夹，用于存放转换后的txt标签文件。运行 2.xml2txt.py 如下代码：

import xml.etree.ElementTree as ET import pickle import os from os import listdir, getcwd from os.path import join import bs4 from PIL import Image   classes = ["person", "hat"] #为了获得cls id def convert(size, box):     dw = 1. / (size[0])     dh = 1. / (size[1])     x = (box[0] + box[1]) / 2.0 - 1     y = (box[2] + box[3]) / 2.0 - 1     w = box[1] - box[0]     h = box[3] - box[2]     x = x * dw     w = w * dw     y = y * dh     h = h * dh     return (x, y, w, h) def convert_annotation(image_id): global none_counts      # 输入文件xml     in_file = open('./labels/valid/%s.xml' % (image_id)) # 输出label txt     out_file = open('./labels/txt_valid/%s.txt' % (image_id), 'w')     tree = ET.parse(in_file)     root = tree.getroot()     size = root.find('size') print(image_id) if size == None: print('{}不存在size字段'.format(image_id)) # 第一个处理方法         img = Image.open('./VOCPerson/JPEGImages/' + image_id + '.jpg')         w, h = img.size  # 大小/尺寸 print('{}.xml缺失size字段, 读取{}图片得到对应 w：{} h：{}'.format(image_id, image_id, w, h)) # # 第二种处理方法 # # 移除xml # os.remove('./VOCPerson/Annotations/' + image_id + '.xml') # # 移除上面被移除掉xml对应的jpg # os.remove('./VOCPerson/JPEGImages/' + image_id + '.jpg')          none_counts += 1 else:          w = int(size.find('width').text)         h = int(size.find('height').text) for obj in root.iter('object'):             cls = obj.find('name').text             if cls not in classes: continue             cls_id = classes.index(cls)             xmlbox = obj.find('bndbox')             b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))             bb = convert((w, h), b)             out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + 'n') if __name__=='__main__':     xml_count = 0     none_counts = 0     list_file = os.listdir('./labels/valid/') for file in list_file: print(file) # image_id = file.replace('.xml', '')         image_id = file.split('.')[0]         convert_annotation(image_id)         xml_count = xml_count + 1 print('没有size字段的xml文件数目：{}'.format(none_counts)) print('总xml个数是 {}'.format(xml_count)) 

得到yolo训练的标签文件。
注意1：这里需要修改文件路径list_file，将train，valid数据集的xml全部转换为txt格式放到对应label文件夹下
注意2：转换完成后将原来的train，valid文件夹删除并将txt_train，txt_valid文件夹重命名看那个为train，valid。

生成train/valid.txt

运行如下代码 3.creattxt.py ，生成train.txt文件和valid.txt文件。train.txt文件内容是train数据集的路径，valid.txt文件内容是valid的数据集路径。train.txt 部分内容如下：

data/images/train/1.jpg data/images/train/10.jpg data/images/train/1000.jpg data/images/train/1001.jpg data/images/train/1002.jpg data/images/train/1005.jpg data/images/train/1006.jpg data/images/train/1007.jpg data/images/train/1008.jpg data/images/train/1009.jpg 

3.creattxt.py 如下：

# 根据训练数据集和验证数据集persontrain.txt and personvalid.txt import os, random, shutil  trainDir = 'images/train/' validDir = 'images/valid/'  train_pathDir = os.listdir(trainDir) # 取图片的原始路径 print('训练集图片数目: {}'.format(len(train_pathDir)))  valid_pathDir = os.listdir(validDir) # 取图片的原始路径 print('验证集图片数目: {}'.format(len(valid_pathDir))) # 删除persontrain.txt and personvalid.txt if(os.path.exists('train.txt')):      os.remove('train.txt') print('删除train.txt成功') if(os.path.exists('valid.txt')):      os.remove('valid.txt') print('删除valid.txt成功') def text_save(root, filename, data): # filename为写入CSV文件的路径，data为要写入数据列表. file = open(filename, 'a') for i in range(len(data)):         s = str(data[i]).replace('[', '').replace(']', '') # 去除[],这两行按数据不同，可以选择         s = 'data/' + root + s.replace("'", '').replace(',', '') + 'n' # 去除单引号，逗号，每行末尾追加换行符 file.write(s) file.close() print("保存文件成功") if __name__ == '__main__':     text_save(trainDir, './train.txt', train_pathDir)     text_save(validDir, './valid.txt', valid_pathDir) print('train.txt 有 {} 行'.format(len([i for i in open('./train.txt', 'r')]))) print('valid.txt 有 {} 行'.format(len([i for i in open('./valid.txt', 'r')]))) 

至此，数据准备工作完成。

YOLOv4配置

需要修改配置的地方主要有三处：cfg文件，data文件，names文件。

cfg文件修改

这里推荐使用yolov4-relu.cfg文件，因为我用了yolov4.cfg训练太吃显存，2080TI的机子batchsize设置为2也会爆显存。其主要原因是mish函数太占显存了。。。yolov4-relu.cfg将mish函数替代为relu函数，大大降低显存使用！

修改其中一处如下，一共有三处，分别对应三个detect header：

[convolutional] size=1 stride=1 pad=1 filters=21 # 修改为 （class数目 + 4 +1） × 3 = 21 activation=linear   [yolo] mask = 6,7,8 anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401 classes=2 #修改为类别数目，这里是2 num=9 jitter=.3 ignore_thresh = .7 truth_thresh = 1 random=1 scale_x_y = 1.05 iou_thresh=0.213 cls_normalizer=1.0 iou_normalizer=0.07 iou_loss=ciou nms_kind=greedynms beta_nms=0.6 

data文件修改

在data/下新建hat.data文件，填入以下内容即可：

classes= 2 train=data/train.txt valid=data/valid.txt names=data/hat.names 

names文件修改

在data/下新建hat.names文件，填入以下类别名字即可：

person hat 

以上，完成配置修改。

模型训练

晚上以上就可以正常开始训练了，train.py文件内容配置修改如下：

if __name__ == '__main__':     parser = argparse.ArgumentParser()     parser.add_argument('--epochs', type=int, default=300) # 500200 batches at bs 16, 117263 COCO images = 273 epochs     parser.add_argument('--batch-size', type=int, default=6) # effective bs = batch_size * accumulate = 16 * 4 = 64     parser.add_argument('--cfg', type=str, default='cfg/yolov4-relu-hat.cfg', help='*.cfg path')     parser.add_argument('--data', type=str, default='data/hat.data', help='*.data path')     parser.add_argument('--multi-scale', action='store_true', help='adjust (67%% - 150%%) img_size every 10 batches')     parser.add_argument('--img-size', nargs='+', type=int, default=[320, 640], help='[min_train, max-train, test]')     parser.add_argument('--rect', action='store_true', help='rectangular training')     parser.add_argument('--resume', action='store_true', help='resume training from last.pt')     parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')     parser.add_argument('--notest', action='store_true', help='only test final epoch')     parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')     parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')     parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')     parser.add_argument('--weights', type=str, default='weights/yolov4.weights', help='initial weights path')     parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')     parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1 or cpu)')     parser.add_argument('--adam', action='store_true', help='use adam optimizer')     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')     opt = parser.parse_args()     opt.weights = last if opt.resume else opt.weights     check_git_status() print(opt) 

训练日志：

电费太贵了！我训练100多点epoch就终止了，验证集map达到89%（之前用YOLOv3验证集map不到60%。。。。），相当可以。

模型测试

放两张YOLOv4的测试图片如下：


可以看到YOLOv4效果还是挺可以。此外我计算YOLOv4的实际推断速度发现，YOLOv4也大大快于YOLOv3，相比YOLOv3又快又准！！真香！！
不愧是怼了一推佐料的大杀器。。。

当然，YOLOv4也有错检漏检的情况，如下图，毕竟是一阶段检测器，但是在精度和准确度的权衡上，YOLOv4是相当可以了！

OVER！