怎么使用YOLO V5训练我们自己的数据集

YOLO v5是YOLO系列中比较好用的一个。在前面的笔记中我们介绍了YOLO v3和YOLO v4的原理，并且从零开始使用pytorch框架重现了YOLO v3的代码。

然而，YOLO v5的算法每天都有在更新，所以我们不再分析YOLO v5的代码。这里我们一起讨论一下YOLO v5的官方代码如何使用，用来训练我们的数据，达到期望的精度和性能。并且掌握其在常见设备上的推理。

前面几篇相关笔记：

1. 使用K-means算法寻找yolo的锚框 - 简书 (jianshu.com)
2. 如何将图片和锚框一起resize——python代码实现 - 简书 (jianshu.com)
   3.YOLO学习笔记——目标检测算法的简介（RCNN, SSD, YOLO) - 简书 (jianshu.com)
   4.YOLOv3网络结构和基础知识学习笔记 - 简书 (jianshu.com)
3. 如何制作YOLOv3模型训练数据集？ - 简书 (jianshu.com)
4. 如何训练YOLOv3模型？pytorch代码实现 - 简书 (jianshu.com)
5. YOLOv3、YOLOv4、YOLOv5之间的区别 - 简书 (jianshu.com)

一、有什么方法能够将模型部署到常见设备上呢？

(1)pytorch官方解决方案

pytorch官方给出了一定的解决方案，能够简便的将模型部署到安卓和IOS的移动设备上。PYTORCH 官方移动设备部署方案官网上有教程，帮助我们部署我们的神经网络。

（2）NCNN腾讯给出的解决方案

ncnn 是腾讯优图实验室首个开源项目，是一个为手机端极致优化的高性能神经网络前向计算框架。ncnn 无第三方依赖，跨平台，手机端 cpu 的速度快于目前所有已知的开源框架。其支持了大部分的NCNN网络。

ncnn 目前已在腾讯多款应用中使用，如 QQ，Qzone，微信，天天 P 图等。

二、使用YOLO v5推理和训练

我们可以从Git Hub上直接将代码打包下载，下面我们给出了链接。
【1】YOLO V5 Git Hub 代码
【2】YOLO v5 的官方中文文档 访问该网站可能需要一些魔法哦

首先，在使用YOLO v5官方源码之前，我们需要确认我们的一些版本信息
python的版本>=3.7, pytorch的版本>=1.7。

在下载代码之后，我们需要安装一些库，我们在YOLO v5根目录下面会找到requirements.txt文件。使用以下代码实现各类库的安装：

pip install -r requirements.txt

这里可能需要几分钟的时间安装。到这里我们就已经做好准备工作啦。

（1）YOLO源码的介绍

YOLO官方提供了五个版本的模型，分别为：n、s、m、l、x

他们的区别在与推理速度和模型精度的不同，例如上表中，从上到下模型的参数量逐渐增加，模型精度逐渐增加，，推理速度随之变慢。我们需要根据项目需求选取不同的模型来训练我们的数据。

如果想使用某个的模型，我们可以在Github上找到上面这个表，每行的第一个表格就是相应的下载链接，选择所需的模型的权重文件进行下载，也可以在代码中下载（第一次使用的时候会自动下载）。权重文件下载之后需要将文件放到根目录下。

源码文件结构：

• train.py：用于训练模型
• detect.py：这里是模型的推理文件，模型训练好之后，我们使用这个文件对图片进行处理。如果训练我们自己的数据还需要对里面的参数进行更改。例如：更改模型权重文件，更改配置文件等
• export.py：将模型导出为自己想要的格式
• data：images是存放检测图片的文件夹，检测的图片将会放到./data/images/目录下；该目录下还有数据集的配置文件。
• runs：每次运行后会生成这个文件夹，保存检测代码。
• classify:用于做图像分类的算法，训练、预测和验证
• models：用于存放模型的配置文件，yolo.py是根据yaml配置文件生成模型的代码，common.py是yolov5中用到的各种模型块。
• segment：用于做语义分割的算法，训练、预测和验证
• utils：各类工具类存放的地方，有激活函数、图像增强、自动批量设置、构建数据集、损失函数等等

（2）如何使用YOLO 推理数据

YOLO v5给出两种方法进行图像推理。

1. 使用pytorch Hub进行推理

import torch

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # or yolov5n - yolov5x6, custom，这里可以选择使用不同的模型

# Images
img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list，这里传入想要识别的图片视频等

# Inference， 将数据输入模型进行推理
results = model(img)

# Results，将结果返回展示
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

2. 使用detect.py进行推理，我们在YOLO v5的官方源码的根目录下能够找到detect.py这个文件，我们运行这个文件就能够实现模型推理。

我们需要在Windows 命令行或者 Linux的终端中进入到YOLO的根目录下输入如下命令，最终结果会保存到run目录中：

python detect.py --weights yolov5s.pt --source 0                               # 我们需要指定模型全中文件
                                               img.jpg                         # 图片路径
                                               vid.mp4                         # 视频路径
                                               screen                          # screenshot
                                               path/                           # directory
                                               list.txt                        # 图片列表文件
                                               list.streams                    # list of streams
                                               'path/*.jpg'                    # glob 
                                               'https://youtu.be/Zgi9g1ksQHc'  # YouTube视频链接
                                               'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream

• --weights：指定模型的权重文件
• --source：指定模型的输入文件（摄像头、网络摄像头、图片、视频、图片列表、网络视频等）选择其中一个即可。

在detect.py文件中修改推理的一些配置。

parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL') # 权重路径
    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)') # 检测图像存储路径
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path') # 数据集配置文件，如果是你自己的数据集还需要修改其内容
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') # 推流图片的高和宽
    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold') # 置信度阈值
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold') # NMS IOU 阈值
    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='show results') # default False，置为True能够边检测边输出图像
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--visualize', action='store_true', help='visualize features')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name') # 输出结果保存路径
    parser.add_argument('--name', default='exp', help='save results to project/name') # 输出结果保存名字
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')

（3）如何使用yolov5进行训练

这里使用的数据集是百度的开源数据集：车辆识别 - 飞桨AI Studio (baidu.com)

第一步，将数据转换为yolo的格式，这里以voc数据格式为例进行转换。

1. 划分训练集、验证集和测试集

"""
这个工具是用于划分数据集的，需要修改数据集的两个路径
"""
import os  
import random 
import argparse 
import tqdm

parse = argparse.ArgumentParser() 
parse.add_argument('--xml_path', default="D:\document\DL\yolov5-master\yolov5-master\data\dataset\car_detect\Annotations", type=str, help="the voc format xml file path")
parse.add_argument('--txt_path', default="D:\document\DL\yolov5-master\yolov5-master\data\dataset\car_detect\ImageSets", type=str, help="output txt file path")
opt = parse.parse_args()

trainval_percent = 0.9 # 训练集和测试集占总数据集的百分数
train_percent = 0.7 # 训练集占上面那个集合的百分数
xmlfilepath = opt.xml_path
txtsavefile = opt.txt_path
total_xml = os.listdir(xmlfilepath)

if not os.path.exists(txtsavefile):
    os.mkdirs(txtsavefile)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
train_val = random.sample(list_index, tv)
train = random.sample(train_val, tr)

file_trainval = open(txtsavefile + "/trainval.txt", "w")
file_test = open(txtsavefile + "/test.txt", "w")
file_train = open(txtsavefile + "/train.txt", "w")
file_val = open(txtsavefile + "/val.txt", "w")

for i in tqdm.tqdm(list_index):
    name = total_xml[i][:-4] + "\n" # 获取文件名
    if i in train_val:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_test.close()
file_train.close()
file_val.close()

2. 将voc格式转换为yolo格式

"""
这个工具适用于将voc格式数据转换为yolo格式的数据,
- 需要修改数据类别变量 
- 修改数据集路径

- 在根目录生成几个文件，为训练集验证集和测试集，文件内是图片路径
- 在labels文件夹下生成多个txt文件，每一行为一个框
- 
"""
import xml.etree.ElementTree as ET
from tqdm import tqdm 
import os  

sets = ["train", "test", "val"]
classes = ["truck", "bus", "SUV", "taxi", "car"]

annotation_path = "D:/document/DL/yolov5-master/yolov5-master/data/dataset/car_detect/Annotations/"
labels_path = "D:/document/DL/yolov5-master/yolov5-master/data/dataset/car_detect/labels/"
imageset_path = "D:/document/DL/yolov5-master/yolov5-master/data/dataset/car_detect/ImageSets/"
root_path = "D:/document/DL/yolov5-master/yolov5-master/data/dataset/car_detect/"
images_path = "D:/document/DL/yolov5-master/yolov5-master/data/dataset/car_detect/images/"

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2. - 1.
    y = (box[2] + box[3]) / 2. - 1.
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw 
    w = w*dw 
    y = y*dh
    h = h*dh
    x = round(x, 6)
    w = round(w, 6)
    y = round(y, 6)
    h = round(h, 6)
    return x,y,w,h

def convertannotation(image_id):
    in_file = open(annotation_path + image_id +".xml", encoding='utf-8')
    out_file = open(labels_path + image_id + ".txt", "w", encoding='utf-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text 
        cls = obj.find('name').text 
        if cls not in classes or int(difficult) == 1:
            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))
        b1, b2, b3, b4 = b 
        # 这里是为了防止标签框超出图片
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h 
        b = (b1, b2, b3, b4)
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + "\n")

wd = os.getcwd()
for image_set in sets:
    if not os.path.exists(labels_path):
        os.makedirs(labels_path)
    img_ids = open(imageset_path + image_set + ".txt").read().strip().split()
    list_file = open(root_path + image_set + ".txt", "w")
    for image_id in tqdm(img_ids):
        list_file.write(images_path + image_id + ".jpg\n")
        convertannotation(image_id)
    list_file.close()

第二步：配置数据集的yaml文件，并将其保存在data目录下

第三步：修改train.py，修改模型文件，修改数据集配置，修改模型配置文件中的分类数量。还需要修改9个锚框。修改utils/dataloaders.py下的img2label_paths函数，将光标处的文件夹名称，改成我们存储图片的文件夹名称。

训练完之后，我们可以在./runs/train/exp/文件夹下面查看自己的训练结果图、代码中最好的模型和最后的模型权重文件（best.pt、last.pt）、可以看到一些看到数据增强的效果图。

如果我们需要推理杠训练好的模型，我们可以设置detect.py里的配置，如前面第二部分提到的图像推理方法一样。

• 主要设置权重文件，将我们的best.pt复制到根目录下，将原来的yolo权重文件替换掉
• 将数据集配置改为自己的数据集yaml配置文件
• 将需要测试的图片、视频放到指定位置

最后运行detect.py，在./runs/detect/exp/里就会出现预测结果。