Yolo v5训练Udacity数据集保姆级教程（Udacity数据集标签转换及数据集划分代码）

一、标签转换为YOLO v5格式

1. 类别处理

        Udacity中的标签类别有Car、Truck、Pedestrain、Biker、Trafficlight 由于博主的研究需要，我们只识别出车辆、行人和骑行者，我们将删去Trafficlight 类，此部分代码如下（注意，本文代码中的文件地址部分均请自行修改）：

# 由于我的项目不检测交通灯，去掉udacity中的交通灯标注
# #以下为去掉udacity数据集中的trafficLight数据的代码

file=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels.txt','r') # 原始labels.txt的地址
newdata = []
for eachline in file:
    data = eachline.strip().split.py(' ')
    if data[6][1:-1] == "trafficLight":
        continue
    else:
        newline = ''
        for i in range(len(data)):
            newline += data[i]
            if i!=len(data)-1:
                newline += ' '
        newline += '\n'
        newdata.append( newline )
file.close()

for d in newdata:
    print(d)
# 不确定写入新文件的数据是否正确时，先不执行以下写newdata到新文件的操作，防止误操作


file1=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels_notraffic.txt','a')
file1.writelines(newdata)
file1.close()

       此外，为和之前的研究统一，我们将Car类和Truck类统一为Car类，Biker类改名为Cyclist，并将所有类别首字母大写，此部分代码如下：

# 统计发现当前txt标注文件只有4类 ['car', 'pedestrian', 'truck', 'biker']
# 为和kitty风格统一，将‘car’和‘pedestrian’合并为‘car’类，以下代码实现

file=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels_notraffic.txt','r')

newdata = []
for line in file:
    data = line.strip().split.py()
    if(data[-1][1:-1] == 'truck'):
        data [-1] = '"car"'
    newline = ''
    for i in range(len(data)):
        newline += data[i]
        if i!=len(data)-1:
            newline += ' '
    newline += '\n'
    newdata.append( newline )
print(newdata)
file.close()

file1=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels_notraffic_truct2car.txt','a')
file1.writelines(newdata)
file1.close()

file=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels_notraffic_truct2car.txt','r')
newdata = []
for line in file:
    data = line.strip().split.py()
    if(data[-1][1:-1] == 'car'):
        data [-1] = '"Car"'
    elif(data[-1][1:-1] == 'biker'):
        print(data[-1],data[-2])
        data [-1] = '"Cyclist"'
        data [-2] = '1'
        print(data[-1], data[-2])
    elif(data[-1][1:-1] == 'pedestrian'):
        print(data[-1], data[-2])
        data[-1] = '"Pedestrian"'
        data [-2] = '2'
        print(data[-1], data[-2])
    newline = ''
    for i in range(len(data)):
        newline += data[i]
        if i!=len(data)-1:
            newline += ' '
    newline += '\n'
    newdata.append( newline )
print(newdata)
file.close()

file1=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels_2Caps.txt','a')
file1.writelines(newdata)
file1.close()

2. 坐标转换与文件分割

        在下载完Udacity数据集后，我们发现其标签格式如下：

         经过选取其中一张图片进行测量，我们确定其格式为：

        PictureName Xmin Ymin Xmax Ymax ClassId ClassName

        而YOLO v5原算法中需要的数据格式为：

        ClassId Midx Midy Width Height

        因此，为避免改动代码，我们选择对标签格式进行转换。ClassId已经有了，我们只需要简单移动其位置即可，当务之急是需要计算出Midx， Midy， Width， Height这几个参数。Midx，Midy为中心点的横纵坐标，其中，Midx计算公式为(Xmin + Ymin) / 2，Midy计算公式为(Ymin + Ymax) / 2 ，由于YOLOv5的标签格式为归一化值，我们还需要对初步计算出来的Midx、Midy进行归一化：Midx /= 图像宽度， Midy /= 图像高度。 对于Width， Height的计算，我们采用公式Width = Xmax - Xmin，Height = Ymax - Ymin，同样，我们需要对其归一化处理：Width /= 图像宽度，Height /= 图像高度。

        同样，为避免改动算法。我们需要将此标签文件进行分割，将同一张图片的标签分割到一个文件中。

        代码如下：

# 由于kitty中坐标是小数，把坐标都变成小数且确定坐标顺序是否一致，删去类别标签名，以下代码实现

file=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/（3）labels_2Caps.txt','r')
# 首先确定udacity中的标签格式为（经过visio测量）：xmin（左边的横坐标，范围从左到右0-1920），ymin（0-1200），xmax，ymax
size_x, size_y = 1920, 1200
# udacity中格式： picture_name, x_min, y_min, x_max, y_max, class_num, class_name
for line in file:
    data = line.strip().split.py()
    picture_name, x_min, y_min, x_max, y_max, class_num, class_name = data[0][:-4], int(data[1]), int(data[2]), int(data[3]), int(data[4]), data[5], data[6]
    txt_path = 'D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels/' + picture_name + '.txt'  # 生成的txt标注文件地址
    txt = open(txt_path, 'a')

    w, h, cx, cy = (x_max - x_min) / size_x, (y_max - y_min) / size_y, (x_min + x_max) / (2 * size_x), (y_min + y_max) / (2 * size_y)
    # 转换公式为 w = (xmax - xmin)/size_x, h = (ymax - ymin)/size_y, cx = (xmin + xmax)/(2*size_x), cy = (ymin + ymax)/(2*size_y)
    w, h, cx, cy = str(w), str(h), str(cx), str(cy)

    newline = ''
    newline += class_num + ' ' + cx + ' ' + cy + ' ' + w + ' ' + h + '\n'
    # 由于yolov5中的标签格式为 类别号， 中心点x坐标cx（范围：从左到右为0到1），中心点y坐标cy（范围：从上到下为0到1），宽度归一化值w，高度归一化值h
    txt.writelines(newline)
    txt.close()

file.close()

二、数据集划分

        代码使用：代码中注释多，文件操作在最后一段，大家可以分段执行，查看打印的信息是否有误，确认无误后执行全部代码进行文件操作。

import os.path
import shutil
import glob
import random

# file=open('D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/（3）labels_2Caps.txt','r')

# 根据 labels 的文件名划分数据集， 不能根据图片， 因为只有交通灯标注的图片没有了标签， 标签数 < 图片数， 只选择有标签的划分数据集

train_percent = 0.7
val_percent = 0.3
num_data = 12996
# 保证随机可复现
random.seed(0)#保证每次随机抽取的都可以复现
# 先划分标签

list_val = []
while len(list_val) < int(num_data * val_percent):
    temp = random.randrange(0, num_data)
    if temp not in list_val:
        list_val.append(temp)
list_train = []
for i in range(num_data):
    if i not in list_val:
        list_train.append(i)
print('验证集数量：', len(list_val), '训练集数量：',  len(list_train),'数量之和：', len(list_val) + len(list_train),'原数据集数量：', num_data)

# 划分好了训练集和验证集的序号，现在按序号找文件名放进一个列表
name_train = []  # 存放无后缀的文件名
name_val = []

labelpath = 'D:/pythonfile/pytorch/datasets/Udacity/object-dataset/object-dataset/labels'
filename = [i.split('.')[0] for i in os.listdir( labelpath )]

for i in range( num_data ):
    if i in list_train:
        name_train.append(filename[i])
    else:
        name_val.append(filename[i])
print('对于文件名列表，验证集数量：', len(name_val), '训练集数量：', len(name_train), '数量之和：', len(name_val) + len(name_train),'原数据集数量：', num_data)

# 现在根据训练集和验证集2个列表的文件名，将对应文件放到对应目录下

rootPath ='D:/pythonfile/pytorch/datasets/Udacity/object-dataset/'

# num_val = 0
# num_train = 0
# for file in glob.glob(rootPath + 'object-dataset/' + '*.jpg'):
#     imagename = file.split('\\')[1].split('.')[0]
#     if imagename in name_train:
#         destpath = rootPath + 'images/train/' + imagename + '.jpg'
#         num_train += 1
#     elif imagename in name_val:
#         destpath = rootPath + 'images/val/' + imagename + '.jpg'
#         num_val += 1
#     shutil.move(file, destpath)
#
# print('划分后的验证集图片数量应有:', num_val, '划分后的训练集图片数量应有', num_train)
num_val = 0
num_train = 0
for file in glob.glob(rootPath + 'images/train/' + '*.jpg'):
    num_train += 1
for file in glob.glob(rootPath + 'images/val/' + '*.jpg'):
    num_val += 1
print('划分后的验证集图片数量实际有:', num_val, '划分后的训练集图片数量实际有', num_train)

num_val = 0
num_train = 0
# 标签划分
for file in glob.glob(rootPath + 'object-dataset/labels/' + '*.txt'):
    labelname = file.split('\\')[1].split('.')[0]
    if labelname in name_train:
        destpath = rootPath + 'labels/train/' + labelname + '.txt'
        num_train += 1
    elif labelname in name_val:
        destpath = rootPath + 'labels/val/' + labelname + '.txt'
        num_val += 1
    shutil.move(file, destpath)
print('划分后的验证集标签应有:',num_val, '划分后的训练集图片数量应有', num_train)