【YOLOv5实战2】基于YOLOv5的交通标志识别系统-自定义数据集

每天进步一点丶

已于 2023-04-26 17:06:23 修改

阅读量3.9k

点赞数 1

分类专栏：深度学习文章标签： python 深度学习 pytorch

于 2022-07-02 23:09:16 首次发布

本文链接：https://blog.csdn.net/a_code_rookie/article/details/125348588

版权

深度学习专栏收录该内容

7 篇文章 34 订阅

订阅专栏

实战博客指引：

项目源代码可联系博主获取。

一、数据准备

1.1 从官网下载YOLOv5

打开官网YOLOv5，使用git或者直接下载压缩包

git下载

git clone git@github.com:ultralytics/yolov5.git

git clone https://github.com/ultralytics/yolov5.git

在这里插入图片描述

直接下载zip压缩包

二、自定义数据集

2.1 创建目录

在yolov5目录下创建文件夹VOCdata，创建图片路径以及标注路径，注意每个标注都对应一个图片名称，并区分训练集与验证集。目录如下：

VOCdata
	--images
		--train
		--val
	--labels
		--train
		--val

2.2 数据预览

下载数据集后，数据集格式如下：
在这里插入图片描述
一共有6列，简要解释一下每列的意义：

第1列：图片名称以及图片格式
第2-5列：图片标注的坐标，以像素为坐标，左上角为原点坐标。
第6列：该标注的分类类别。

2.3 数据格式转换

2.3.1 图像格式转换

首先将ppm格式转换为jpg格式，在VOCdata/下目录下创建py文件ppm_to_jpg.py:

def ppm_to_jpg(img_out_path, ppm_path):
    ppm_list = os.listdir(ppm_path)
    for ppm in ppm_list:
        if ppm.endswith(".ppm"):
            path = ppm_path + ppm
            img = Image.open(path)
            new_path = img_out_path + ppm.split(".ppm")[0] + ".jpg"
            img.save(new_path)
    print("transfer finish!")

首先指定ppm图片的输入路径以及输出路径，并创建文件夹，用于保存训练集和测试集。注意测试机保存在yolov5/data/images/test/目录下，用于后面进行数据集的测试。

2.3.2 标注格式转换

在YOLOv5中，标注坐标有要求格式，因此需要进行转换。官方的原文如下：

然后准备labels，将xml文件转换成yolo需要的txt格式。
After using a tool like CVAT, makesense.ai or Labelbox to label your images, export your labels to YOLO format, with one *.txt file per image (if no objects in image, no *.txt file is required). The *.txt file specifications are:

- One row per object
- Each row is class x_center y_center width height format.
- Box coordinates must be in normalized xywh format (from 0 - 1). If your boxes are in pixels, divide x_center and width by image width, and y_center and height by image height.
- Class numbers are zero-indexed (start from 0).

在这里插入图片描述
The label file corresponding to the above image contains 2 persons (class 0) and a tie (class 27):

接下来进行标注转换。新建py文件split_labels.py，用于将标注文件进行格式转换并划分数据集。这里采用训练集：验证集=9：1，在后续中的图片划分中将按照该labels文件夹划分图像数据集。

def gt2yolo():
    i = 0
    file_name = ''
    for line in open(input_path, "r"):  # 设置文件对象
        words = line.split(';')
        img_name = words[0].split(".")[0]
        label = int(words[-1])
        x = float(words[1]) / orign_w
        w = float(words[3]) / orign_w - x
        y = float(words[2]) / orign_h
        h = float(words[4]) / orign_h - y
        # 保存路径。并设置是保存为训练集还是验证集
        if output_path + img_name + ".txt" != file_name:
            if i < int(nums * 0.9):
                file_name = output_path + 'train/' + img_name + ".txt"
            else:
                file_name = output_path + 'val/' + img_name + ".txt"
            fw = open(file_name, 'w')
        else:
            fw = open(file_name, 'a')
        # print(file_name)
        # 写入文件中
        fw.write(str(label) + " " + str(x) + " " + str(y) + " " + str(w) + " " + str(h) + "\n")
        fw.close()
        i += 1

再看一遍标注文件的格式:
在这里插入图片描述
因此基本思路为按照";"划分数据，并将中间四个像素坐标归一化到[0,1]区间上。详细过程在代码中写的很详细，注意要修改自己的输入输出路径以及图片的大小。
划分完成后，将得到如下的目录结构：

2.3.3 划分图片

划分好标注数据集后，将按照该数据集划分图像数据。基本思路为遍历训练集和验证集，获取对应标注的文件名称，然后找到相应的文件名称后将图片移动至对应的训练集或验证集中。详细注释见代码，注意修改路径。
创建split_train_val.py：

def train_val_img_list(img_labels, mode):
    # 获取标注路径下的图片名称
    img_list = os.listdir(os.path.join(img_labels, mode))
    # print(img_list)
    # 切片，只需要名称，不需要格式
    return [i[:-4] for i in img_list]


def split_train_val():
    # 遍历训练集与验证集标注目录，获取图片名称
    for i in label_list:
        img_list = train_val_img_list(img_labels, i)
        # 指定输入位置与保存位置
        for img_name in img_list:
            # D:/download/Train_jpg/00000.jpg
            in_img_path = os.path.join(input_path, str(img_name)) + ".jpg"
            # D:/yolov5/VOCdata/images/train\00000.jpg
            out_img_path = os.path.join(out_path, i, str(img_name)) + '.jpg'
            shutil.copy(in_img_path, out_img_path)

        print("copy finish!!")