yolov5使用自己的数据集训练

一、环境与配置

1.1 安装anaconda并创建虚拟环境

参考 安装anaconda并创建虚拟环境

1.2 安装CUDA和cudnn

参考 安装CUDA 和cudnn

1.3 安装pytorch

参考 安装pytorch

二、源码测试

2.1 下载源码

在github上下载源码：yolov5源码下载地址

2.2 安装依赖项

下载源码并解压后，源码根目录下有一个requirements.txt，这里面就是需要安装的各种依赖项了，安装方法是，从根目录打开命令提示符：在文件夹上方的框里输入“cmd”，然后按回车：

输入命令：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

更换镜像源，可以提高下载速度。

2.3 测试

2.3.1 下载权重文件

作者在GitHub中给出了他们训练出来的权重文件：
yolov5权重文件下载地址
并给出了各种权重文件的检测效果，我们可以随意下载。
将下载的权重文件放在./weight文件夹下。

2.3.2 测试

还是在源码的根目录下打开cmd，然后激活pytorch环境。

图片测试
输入命令：

python detect.py --source ./data/images/ --weights ./weights/yolov5x6.pt --conf 0.4

打开根目录下多出来的runs文件夹，可以看到测试结果。
在这里插入图片描述

视频测试
准备一个视频文件test.mp4，放在/data/video路径下，输入命令：

python detect.py --source ./data/video/test.mp4 --weights ./weights/yolov5x6.pt --conf 0.4

三、训练自己的数据集

3.1 数据集的制作

参考这里制作VOC数据集。
第一步：将之做好的数据集放到yolov5-master根目录下，数据集文件夹结构为：

第二步：将数据集转换到yolo数据集格式。在VOCdata中新建一个voc_label.py文件，输入如下代码：

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

sets=[('TrainVal', 'train'), ('TrainVal', 'val'), ('Test', 'test')]

classes = ['类别1', '类别2', '类别3', '类别4']

def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    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(year, image_set, image_id):
    in_file = open('VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOC%s/labels/%s_%s/%s.txt'%(year, year, image_set, image_id), 'w')
    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))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

def copy_images(year,image_set, image_id):
    in_file = 'VOC%s/JPEGImages/%s.jpg'%(year, image_id)
    out_flie = 'VOC%s/images/%s_%s/%s.jpg'%(year, year, image_set, image_id)
    shutil.copy(in_file, out_flie)

wd = getcwd()

for year, image_set in sets:
    if not os.path.exists('VOC%s/labels/%s_%s'%(year,year, image_set)):
        os.makedirs('VOC%s/labels/%s_%s'%(year,year, image_set))
    if not os.path.exists('VOC%s/images/%s_%s'%(year,year, image_set)):
        os.makedirs('VOC%s/images/%s_%s'%(year,year, image_set))
    image_ids = open('VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('VOC%s/%s_%s.txt'%(year, year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOC%s/images/%s_%s/%s.jpg\n'%(wd, year, year, image_set, image_id))
        convert_annotation(year, image_set, image_id)
        copy_images(year, image_set, image_id)

    list_file.close()

转换后可以看到VOCData/VOCTrainval/和VOCData/VOCTest下生成了两个新的文件夹images和labels，以及相应的训练、验证、或测试的txt文件。


其中images和label文件夹下又有以相应的txt文件名命名的文件夹，里面存放了对应txt文件内容中的图片和标签信息。
现在整个数据集文件夹的结构如下：

3.2 修改配置文件

3.2.1 修改数据集方面的yaml文件

在准备完数据后，v5创新性的省略了.data和.names文件的配置，而是将二者合二为一到yaml中，在data文件夹下创建myvoc.yaml文件，输入以下信息

3.2.2 聚类获得先验框

在VOCData目录下创建程序两个程序 kmeans.py 以及 clauculate_anchors.py

不需要运行 kmeans.py，运行 clauculate_anchors.py 即可。

kmeans.py 程序如下：这不需要运行，也不需要更改

import numpy as np

def iou(box, clusters):
    """
    Calculates the Intersection over Union (IoU) between a box and k clusters.
    :param box: tuple or array, shifted to the origin (i. e. width and height)
    :param clusters: numpy array of shape (k, 2) where k is the number of clusters
    :return: numpy array of shape (k, 0) where k is the number of clusters
    """
    x = np.minimum(clusters[:, 0], box[0])
    y = np.minimum(clusters[:, 1], box[1])
    if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:
        raise ValueError("Box has no area")    # 如果报这个错，可以把这行改成pass即可

    intersection = x * y
    box_area = box[0] * box[1]
    cluster_area = clusters[:, 0] * clusters[:, 1]

    iou_ = intersection / (box_area + cluster_area - intersection)

    return iou_

def avg_iou(boxes, clusters):
    """
    Calculates the average Intersection over Union (IoU) between a numpy array of boxes and k clusters.
    :param boxes: numpy array of shape (r, 2), where r is the number of rows
    :param clusters: numpy array of shape (k, 2) where k is the number of clusters
    :return: average IoU as a single float
    """
    return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])

def translate_boxes(boxes):
    """
    Translates all the boxes to the origin.
    :param boxes: numpy array of shape (r, 4)
    :return: numpy array of shape (r, 2)
    """
    new_boxes = boxes.copy()
    for row in range(new_boxes.shape[0]):
        new_boxes[row][2] = np.abs(new_boxes[row][2] - new_boxes[row][0])
        new_boxes[row][3] = np.abs(new_boxes[row][3] - new_boxes[row][1])
    return np.delete(new_boxes, [0, 1], axis=1)


def kmeans(boxes, k, dist=np.median):
    """
    Calculates k-means clustering with the Intersection over Union (IoU) metric.
    :param boxes: numpy array of shape (r, 2), where r is the number of rows
    :param k: number of clusters
    :param dist: distance function
    :return: numpy array of shape (k, 2)
    """
    rows = boxes.shape[0]

    distances = np.empty((rows, k))
    last_clusters = np.zeros((rows,))

    np.random.seed()

    # the Forgy method will fail if the whole array contains the same rows
    clusters = boxes[np.random.choice(rows, k, replace=False)]

    while True:
        for row in range(rows):
            distances[row] = 1 - iou(boxes[row], clusters)

        nearest_clusters = np.argmin(distances, axis=1)

        if (last_clusters == nearest_clusters).all():
            break

        for cluster in range(k):
            clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)

        last_clusters = nearest_clusters

    return clusters

if __name__ == '__main__':
    a = np.array([[1, 2, 3, 4], [5, 7, 6, 8]])
    print(translate_boxes(a))

运行：clauculate_anchors.py

会调用 kmeans.py 聚类生成新anchors的文件

程序如下：

需要更改第 9 、13行文件路径 以及 第 16 行标注类别名称

# -*- coding: utf-8 -*-
# 根据标签文件求先验框

import os
import numpy as np
import xml.etree.cElementTree as et
from kmeans import kmeans, avg_iou

FILE_ROOT = "D:/Yolov5/yolov5/VOCData/"     # 根路径
ANNOTATION_ROOT = "Annotations"   # 数据集标签文件夹路径
ANNOTATION_PATH = FILE_ROOT + ANNOTATION_ROOT

ANCHORS_TXT_PATH = "D:/Yolov5/yolov5/VOCData/anchors.txt"   #anchors文件保存位置

CLUSTERS = 9
CLASS_NAMES = ['light', 'post']   #类别名称

def load_data(anno_dir, class_names):
    xml_names = os.listdir(anno_dir)
    boxes = []
    for xml_name in xml_names:
        xml_pth = os.path.join(anno_dir, xml_name)
        tree = et.parse(xml_pth)

        width = float(tree.findtext("./size/width"))
        height = float(tree.findtext("./size/height"))

        for obj in tree.findall("./object"):
            cls_name = obj.findtext("name")
            if cls_name in class_names:
                xmin = float(obj.findtext("bndbox/xmin")) / width
                ymin = float(obj.findtext("bndbox/ymin")) / height
                xmax = float(obj.findtext("bndbox/xmax")) / width
                ymax = float(obj.findtext("bndbox/ymax")) / height

                box = [xmax - xmin, ymax - ymin]
                boxes.append(box)
            else:
                continue
    return np.array(boxes)

if __name__ == '__main__':

    anchors_txt = open(ANCHORS_TXT_PATH, "w")

    train_boxes = load_data(ANNOTATION_PATH, CLASS_NAMES)
    count = 1
    best_accuracy = 0
    best_anchors = []
    best_ratios = []

    for i in range(10):      ##### 可以修改，不要太大，否则时间很长
        anchors_tmp = []
        clusters = kmeans(train_boxes, k=CLUSTERS)
        idx = clusters[:, 0].argsort()
        clusters = clusters[idx]
        # print(clusters)

        for j in range(CLUSTERS):
            anchor = [round(clusters[j][0] * 640, 2), round(clusters[j][1] * 640, 2)]
            anchors_tmp.append(anchor)
            print(f"Anchors:{anchor}")

        temp_accuracy = avg_iou(train_boxes, clusters) * 100
        print("Train_Accuracy:{:.2f}%".format(temp_accuracy))

        ratios = np.around(clusters[:, 0] / clusters[:, 1], decimals=2).tolist()
        ratios.sort()
        print("Ratios:{}".format(ratios))
        print(20 * "*" + " {} ".format(count) + 20 * "*")

        count += 1

        if temp_accuracy > best_accuracy:
            best_accuracy = temp_accuracy
            best_anchors = anchors_tmp
            best_ratios = ratios

    anchors_txt.write("Best Accuracy = " + str(round(best_accuracy, 2)) + '%' + "\r\n")
    anchors_txt.write("Best Anchors = " + str(best_anchors) + "\r\n")
    anchors_txt.write("Best Ratios = " + str(best_ratios))
    anchors_txt.close()

会生成anchors文件。如果生成文件为空，重新运行即可。

第二行 Best Anchors 后面需要用到。

3.2.3 修改网络参数方面的yaml文件

这个相当于以前版本的.cfg文件，在models/yolov5x.yaml【当然，你想用哪个模型就去修改对应的yaml文件】，就修改一下类别数量;
修改anchors，根据 anchors.txt 中的 Best Anchors 修改，需要取整（四舍五入、向上、向下都可以）。
保持yaml中的anchors格式不变，按顺序一对一即可。

3.2.4 修改train.py中的一些参数

在train.py中修改一下训练参数，也可以直接在训练语句中重写，修改的话只是修改默认值。

parser.add_argument('--epochs', type=int, default=200)  # 根据需要自行调节训练的epoch
parser.add_argument('--batch-size', type=int, default=8)  # 根据自己的显卡调节，显卡不好的话，就调小点
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='*.cfg path') # 根据需要，自行选择模型
parser.add_argument('--data', type=str, default='data/myvoc.yaml', help='*.data path')  # data设置为前两步中我们新建的myvoc.yaml
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes')  # 可调可不调
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  # 使用CPU还是GPU训练

3.3 训练

3.3.1 在预训练模型基础上进行迁移训练

在yolov5-master目录中打开命令提示符，再激活torch环境，然后输入下列命令：

python train.py --img 640 --batch 8 --epoch 200 --data ./data/myvoc.yaml --cfg ./models/yolov5x.yaml --weights ./weights/yolov5x6.pt --device 0 --workers 8

参数解释：

--img：输入图片尺寸
--batch：batch_size大小
--epoch：训练周期
--data：上面修改过的数据集方面的信息文件
--cfg：模型的配置文件，自行选择模型，当然这里选的是哪个模型，就要在3.2.2中修改相应的配置文件中的类别数
--weights：预训练模型，之前下载好的，这里选择的预训练模型需要跟–cfg中相同
--device：CUDA设备，如果使用CPU训练就写cpu，如果用一块GPU训练就写0，如果有多块GPU就写 0,1,2,3
--workers：线程数

其他可修改参数：

--rect：rectangular training，布尔值
--resume：是否从最新的last.pt中恢复训练，布尔值
--nosave：仅仅保存最后的checkpoint，布尔值
--notest：仅仅在最后的epoch上测试，布尔值
--evolve：进化超参数（evolve hyperparameters），布尔值
--bucket：gsutil bucket，默认值’’
--cache-images：缓存图片可以更快的开始训练，布尔值
--name：如果提供，将results.txt重命名为results_name.txt
--adam：使用adam优化器，布尔值
--multi-scale：改变图片尺寸img-size +/0- 50%，布尔值
--single-cls：训练单个类别的数据集，布尔值

当然这些参数可以直接在train.py文件中直接修改默认值，然后直接运行train.py文件。

3.3.2 从头开始训练

将train.py中的第458行–weights的默认参数删掉

在上面命令的基础上，去掉–weight，在yolov5-master目录中打开命令提示符，再激活torch环境，然后输入下列命令：

python train.py --img 640 --batch 8 --epoch 200 --data ./data/myvoc.yaml --cfg ./models/yolov5x.yaml --device 0 --workers 8

当然如果其他参数也修改为适合自己训练的参数了，也可以直接运行train.py文件。

3.3.3 tensorboard可视化训练过程

在yolov5-master目录中打开命令提示符，再激活torch环境，然后输入下列命令：

tensorboard --logdir=runs/train/exp

参数解释：
tensorboard：可视化tensorboard命令
–logdir：训练日志文件所在目录（通常在训练的时候会自动建立runs/train/expn这么一个路径，其训练过程中的所有信息都保存在这个文件夹下，包括训练的日志文件。其中expn不是一个固定的文件夹名字，这里面的n是根据训练的次数依次进行排序的，在进行可视化的时候注意找对文件夹）
然后复制下面的地址，用任意一个浏览器打开

得到训练过程中的各种指标的变化信息：

可能遇到的问题：
若在浏览器中无法打开改地址，可能的原因是6006端口号被占用，这时候可以指定从其他端口打开，命令如下：

tensorboard --logdir=runs/train/exp --port=6008

3.4 测试

3.4.1 计算mAP

主要使用根目录下的test.py文件，有两种方法，建议使用方法二，因为test.py文件在训练过程中也会使用到，第一种方法修改了默认的参数，可能会导致以后再次训练时出错。
方法一：修改test.py文件中的相关参数，主要是修改–weights模型地址、–data数据集配置文件地址，–task当前进行的任务，这三个参数，其他的都可以看情况自行修改。

各种可修改参数的解释：

--weights ：模型路径，
--data：数据集的配置文件
--batch-size：默认值32
--img-size：图片大小，默认640
--conf-thres：目标置信度阈值
--iou-thres：NMS的IOU阈值
--save-json：把结果保存为cocoapi-compatible的json文件
--task：当前进行的任务，改为test，可选其他值：val, test, study
--device：cuda设备，例如：0或0,1,2,3或cpu，默认’’
--half：半精度的FP16推理
--single-cls：将其视为单类别，布尔值
--augment：增强推理，布尔值
--verbose：显示类别的mAP，布尔值

修改完之后，直接运行test.py就能够得到测试的结果了，输出内容：

并且在runs\test\exp目录下有更多测试的细节信息：

使用该方法进行测试后，在下一次训练之前一定要将–task参数的默认值改回val
方法二：使用命令给定各个参数值进行测试，不修改test.py文件中的参数。
在yolov5-master目录中打开命令提示符，再激活torch环境，然后输入下列命令：

python test.py --weights ./runs/train/exp/weights/best.pt --data ./data/myvoc.yaml --batch 32 --img 640 --task test --device 0 

其他可修改的参数和运行结果同上。

3.4.2 图片或视频测试

同样也有两种方式可以
在yolov5-master目录中打开命令提示符，再激活torch环境，然后输入下列命令：

python detect.py --source ./data/images/ --weights ./runs/train/exp14/weights/best.pt --conf 0.4

参数解释：

--weights ：模型路径，根据自己模型的保存位置填写
--source：输入的数据源，可以是：图片、目录、视频、网络摄像头、http和rtsp流。
--output： 输出检测结果的路径，默认值为：inference/output
--img-size ：图片的大小（pixels），默认值为：640
--conf-thres：对象的置信度阈值（object confidence threshold），默认值为：0.4
--iou-thres ：NMS的IOU阈值（ IOU threshold for NMS），默认值为：0.5
--device：cuda设备，例如：0或0,1,2,3或cpu，默认’’
--view-img ：显示结果，‘布尔值，默认为true’
--save-txt ：把结果保存到*.txt文件中
--classes：过滤类别 CLASSES [CLASSES …]，filter by class
--agnostic-nms：类不可知 NMS
--augment：增强推理（augmented inference）

对–source参数的特别说明：
（1）检测单张图片，直接写该图片的路径加上图片名，如–source ./data/images/test.jpg；
（2）检测整个文件夹下所有的图片，就不用具体些到图片名，路径写到该文件夹就成，例如：–source ./data/images/，表示检测images文件夹下所有的图片。
（3）检测视频文件：直接写路径加视频的名称，如：–source ./data/images/test.mp4
（4）网络摄像头传输视频，使用电脑自带的摄像头检测：将source参数设置为0–source 0（我使用的台式机，没有自带的摄像头，暂时还没测试这一条）
（5）rtsp传输视频，使用外部摄像头，source参数设置为：–source + rtsp连接命令，不同厂家的摄像头使用rtsp连接的方式有些许差别，如海康威视的摄像头连接命令为–source rtsp://admin:密码@IP地址:554/STD_H265/ch1/main/av_stream，其中验证码为摄像头上贴的验证码，一个设备一个验证码；IP地址为当前摄像头的IP地址，摄像头通过网络进行连接，将会被分配一个IP，通过海康威视的官方APP能够查看该设备的IP。

参考（感谢）
https://blog.csdn.net/qq_40927867/article/details/115768888
https://blog.csdn.net/qq_45945548/article/details/121701492