yolov5训练、推理自己的数据集

使用的版本为：yolov5-6.1

1、准备数据集

1.1 图像标注

使用labelimg或者CVAT等标注工具对图像样本进行逐一标注，标注格式选用YOLO（减少不必要的转格式麻烦），每一个图像文件xxxx.jpg标注完成后，都会生成对应的label文件xxxx.txt。labelimg标注示意：

1.2 数据集目录结构

最终的数据集目录结构采取如下方式：

--cat-dog-dataset
	--images
		--train
			cat_1.jpg
			cat_3.jpg
			...
		--val
			cat_2.jpg
			cat_19.jpg
			...
		--test
			cat_13.jpg
			cat_34.jpg
			...
	--labels
		--train
			cat_1.txt
			cat_3.txt
			...
		--val
			cat_2.txt
			cat_19.txt
			...
		--test
			cat_13.txt
			cat_34.txt
			...

数据集划分参考代码如下：

import os
import shutil
import random
random.seed(0)
def split_data(img_path, txt_path, new_file_path, train_rate, val_rate, test_rate):
    directory = ["/images", "/labels", "/train", "/val", "/test"]
    images_dir = new_file_path + directory[0]
    labels_dir = new_file_path + directory[1]
    if not os.path.exists(images_dir):
        os.makedirs(images_dir)

    if not os.path.exists(labels_dir):
        os.makedirs(labels_dir)

    for i in range(2, 5):
        if not os.path.exists(images_dir + directory[i]):
            os.makedirs(images_dir + directory[i])
        if not os.path.exists(labels_dir + directory[i]):
            os.makedirs(labels_dir + directory[i])

    all_images = os.listdir(img_path)
    all_labels = os.listdir(txt_path)
    data = list(zip(all_images, all_labels))
    random.shuffle(data)
    shuffle_images, shuffle_labels = zip(*data)
    train_images = shuffle_images[0: int(len(all_images) * train_rate)]
    train_labels = shuffle_labels[0: int(len(all_labels) * train_rate)]
    val_images = shuffle_images[int(len(all_images) * train_rate): int(len(all_images) * (train_rate + val_rate))]
    val_labels = shuffle_labels[int(len(all_labels) * train_rate): int(len(all_labels) * (train_rate + val_rate))]
    test_images = shuffle_images[int(len(all_images) * (train_rate + val_rate)):]
    test_labels = shuffle_labels[int(len(all_labels) * (train_rate + val_rate)):]

    for file in train_images:
        file_prefix = file.split(".")[0]
        origin_img_path = img_path + "/" + file_prefix + ".jpg"
        new_img_path = images_dir + directory[2] + "/" + file_prefix + ".jpg"
        origin_label_path = txt_path + "/" + file_prefix + ".txt"
        new_label_path = labels_dir + directory[2] + "/" + file_prefix + ".txt"
        shutil.copy(origin_img_path, new_img_path)
        shutil.copy(origin_label_path, new_label_path)

    for file in val_images:
        file_prefix = file.split(".")[0]
        origin_img_path = img_path + "/" + file_prefix + ".jpg"
        new_img_path = images_dir + directory[3] + "/" + file_prefix + ".jpg"
        origin_label_path = txt_path + "/" + file_prefix + ".txt"
        new_label_path = labels_dir + directory[3] + "/" + file_prefix + ".txt"
        shutil.copy(origin_img_path, new_img_path)
        shutil.copy(origin_label_path, new_label_path)

    for file in test_images:
        file_prefix = file.split(".")[0]
        origin_img_path = img_path + "/" + file_prefix + ".jpg"
        new_img_path = images_dir + directory[4] + "/" + file_prefix + ".jpg"
        origin_label_path = txt_path + "/" + file_prefix + ".txt"
        new_label_path = labels_dir + directory[4] + "/" + file_prefix + ".txt"
        shutil.copy(origin_img_path, new_img_path)
        shutil.copy(origin_label_path, new_label_path)

if __name__ == '__main__':
    img_path = "/home/temp/images"
    txt_path = "/home/temp/labels"
    new_file_path = "/home/dataset/cat-dog-dataset"
    split_data(img_path, txt_path, new_file_path, 0.8, 0.1, 0.1)

以上代码实现了将数据集按照8:1:1的比例划分为训练集、验证集和测试集。

2、训练

2.1 创建数据集配置文件

yolov5-6.1源代码data目录下创建文件mydataset.yaml，并将data目录下的coco128.yaml文件内容粘贴进去，然后修改相应内容，如下：

# mydataset.yaml
path: /home/dataset/cat-dog-dataset  # dataset root dir
train: images/train  # train images(relative to 'path')
val: images/val  # val images(relative to 'path')
test:  # test images (optional)

# Classes
nc: 2  # number of classes
names: ['cat', 'dog']  # class names

2.2 修改模型配置文件

在models目录下选择一个yolov5模型，本人选择的是yolov5m，将其中的检测类别数nc修改为正确数值：

# yolov5m.yaml
nc: 80  # number of classes
depth_multiple: 0.67  # model depth multiple
width_multiple: 0.75  # layer channel multiple

2.3 下载预训练权重

yolov5m6.pt的预训练权重下载地址：yolov5m6_local.pt，提取密码：ge2y
将其下载后，放入yolov5-6.1目录下。

2.4 执行命令开始训练

python train.py --cfg models/yolov5m.yaml --data data/mydataset.yaml --weights yolov5m6_local.pt --epochs 100 --device 0,1,2

3、测试

python detect.py --weights runs/train/exp/weights/best.pt --source data/images --data data/mydataset.yaml --device 0