用YOLOv5训练自己的目标检测数据集（以口罩检测为例）

        运行YOLOv5首先需要安装深度学习环境，教程请看安装pytorch深度学习环境（GPU版）。

        YOLOv5的代码在GitHub上是开源的GitHub - ultralytics/yolov5，利用其代码实现自己的目标检测需求，需要3个步骤：1.准备数据集；2.配置代码参数，训练模型；3.预测。以下笔者将带大家一步步实现自己的目标检测模型训练。

一、准备数据集

1.1 收集图片

        我们根据自己的需求收集相关图片，在这里以口罩识别为例。我们从网上收集到一些戴口罩和不戴口罩的图片，如下图所示：                

1.2 利用labelimg软件给收集到的图片打标签

1.2.1 labelimg软件的安装

        labelimg软件是一款开源的数据标注工具，可以标注三种格式。①VOC标签格式的xml文件。②yolo标签格式的txt文件。③createML标签格式的json文件。

        labelimg的安装很简单，我们打开cmd，输入以下命令即可：

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

1.2.2 利用labelimg软件打标签 

        首先我们不妨建立个名为VOC2007的文件夹，里面创建一个名为JPEGImages的文件夹用以存放我们收集好的需要打标签的图片；再创建一个名为Annotations的文件夹用以存放标注的标签文件；最后创建一个名为 predefined_classes.txt 的txt文件用以存放所要标注的类别名称。结构如下图所示：

        在这里我们想要实现的是检测有没有戴口罩，因此predefined_classes.txt文件的类别只有2个，如下图所示：

                                   

        然后，我们要在VOC2007的目录下(一定要是该目录下)打开cmd,输入以下命令： 

labelimg JPEGImages predefined_classes.txt

         这个命令的意思是利用labelimg软件给JPEGImages文件夹中的图片，按照 predefined_classes.txt 文件中的分类打标签。

        打开后的界面如下图所示，其中

        Open Dir是选择存放图片的文件夹，在这里我们的命令将其默认为JPEGImages文件夹；

        Change Save Dir是改变存储标签的文件夹，这里我们默认为Annotations文件夹；

        PascalVOC是选择标签格式，上边介绍过，主要有3种，我们通常选择PascalVOC的xml格式，YOLO格式也行，两者之间可以相互转换；

        Create RectBox是产生打标签的十字位置线，对图片进行标注。

        框取目标检测位置后会出现标签选择框，我们选择对应的标签即可，如下图所示。然后就可以点击Next Image对下一幅图片进行标注，直至全部图片标注完成。

        两种标签格式如下图所示：

        PascalVOC的xml格式：

                                           

        YOLO的txt格式：

                                     

1.3 标签格式的转化及训练集和验证集的划分

1.3.1 xml格式标签转化为txt格式，并划分训练集（80%）和验证集（20%）

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile

classes = ["unmask", "mask"]

TRAIN_RATIO = 80        %训练集的比例


def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / siz