划分自己的数据集

0 说明

最近正在使用yolov4来进行安全帽是否佩戴的检测。但是因为数据集是由不同的小组成员标定的，在标定的过程中有些相似的数据集虽然没有标定，但是也忘记删除了，这就导致了标签个数和数据集的个数不相符的现象出现（图片9608个，标签9263个）。因此，为了更方便的将数据集划分为训练集和测试集，我打算将没有进行标定的数据集进行删除，然后重新划分为训练集和测试集。

 

1 删除未标定的图片

1.1 明确数据和标签的个数

# 获取指定文件夹下指定类型的数量
def getFilsNum(path_dir, type_files):
    # step0: 判断文件夹是否存在，或者是否是个目录
    # step1: 打开文件夹
    # step2: 获取指定文件类型的个数
    # step3: 返回个数
    if not os.path.exists(path_dir) or not os.path.isdir(path_dir):
        print("{0}文件夹不存在，或者不是一个目录".format(path_dir))
        return 0

    num = 0
    files_name = []
    for path_file in os.listdir(path_dir):
        file_name, postfix = path_file.split('.')
        if postfix == type_files:
            num += 1
            files_name.append(file_name)
    return files_name, num


# step1: 打开两个文件夹
path_images_dir = r'D:\Dataset\helmet\helmetAll2020\JPEGImages'
path_annotations_dir = r'D:\Dataset\helmet\helmetAll2020\Annotations'


# files_data和files_lable用来比较哪些是多余的图片或者标签
files_data,length_imgs = getFilsNum(path_images_dir,'jpg')
files_label,length_ans = getFilsNum(path_annotations_dir,'xml')

打印运行的结果：

print(length_ans)
# 9263
print(length_imgs)
# 9608

小结：在这里小结一下使用到的函数

（1）os.listdir()：返回某个目录下面所有的文件（带后缀，你可以使用split进行分割提取）

（2）os.walk()：这个和listdir()功能相似。它的返回值有三个root,dirs,files。root表示当前文件夹的地址，dirs是一个列表，其中保存当前文件夹下所有的目录；files也是一个列表，其中保存当前文件夹下所有的文件

1.2 删除多余的图片或者标签

# 如果图片的数量和标签的数量不相等就删除多余的图片或者标签（在这边图片的个数多余标签的个数，要删除多余的图片）
redun_files = []
# 如果图片的数量多
if length_imgs > length_ans:
    for file_ in files_data:
        if file_ not in files_label:
            redun_files.append(file_)
    # 打开图片目录，获取所有文件名称，如果文件名称在redun_files中，就打印
    for path_file in os.listdir(path_images_dir):
        file_name,_ = path_file.split('.')
        if file_name in redun_files:
            # 删除图片
            os.remove(os.path.join(path_images_dir,path_file))

else: # 标签的数量比图片的多
    for file_ in files_label:
        if file_ not in files_data:
            redun_files.append(file_)
        if file_name in redun_files:
            # 删除标签
            os.remove(os.path.join(path_images_dir,path_file))

print(len(redun_files))
# 345

 

2 划分数据集

2.1 创建训练集和测试集的目录

训练集和测试集的目录我是这样安排的：

明确目录结构之后，我们就需要创建相关的目录

# 创建相关的目录
path_datasets_train = r'D:\Dataset\helmet\cumtHelmet2020\datasets\train'
path_datasets_val = r'D:\Dataset\helmet\cumtHelmet2020\datasets\val'
path_labels_train = r'D:\Dataset\helmet\cumtHelmet2020\labels\train'
path_labels_val = r'D:\Dataset\helmet\cumtHelmet2020\labels\val'
makedirs(path_datasets_train)
makedirs(path_datasets_val)
makedirs(path_labels_train)
makedirs(path_labels_val)

2.2 划分数据集

在这里我是这样处理的，先移动20%验证集的图片和标签，剩下的部分自然就是训练集的图片和标签了，然后再一次性将训练图片和标签移动到指定位置

# 按照8：2的比例划分数据集（先移动验证集，然后将剩余的图片一次性移动到指定位置）
import random
import shutil
def partition_datasets(num_files_in_dir, path_dest, path_source,val_rate):
    files_source = os.listdir(path_source)
    if val_rate>0.5:
        print('请设置验证集的比重')
        return -1
    # 获得要移动图片的数量
    num_part_imgs = int(val_rate*num_files_in_dir)
    # sample一定要是个类别
    sample = random.sample(files_source,num_part_imgs)
    imgs_sample_list = []
    for name in sample:
        # print(os.path.join(path_source,name))
        shutil.move(os.path.join(path_source,name),os.path.join(path_dest,name))
        imgs_sample_list.append(name.split('.')[0])
    return imgs_sample_list

imgs_list = partition_datasets(length_imgs,path_dest=path_datasets_val,path_source=path_images_dir,val_rate=0.2)

# 开始移动相关的标签（先移动验证集的label）
def partitioin_anno(imgs_list,path_dest, path_source):
    num = 0
    for path_label in os.listdir(path_source):
        label,_ = path_label.split('.')
        if label in imgs_list:
            # 移动标签到指定位置
            shutil.move(os.path.join(path_source,path_label),os.path.join(path_dest,path_label))
            # print(os.path.join(path_dest,path_label))

小结：

（1）random.sample()：是用来随机抽取样本的。它的第一个参数是一个列表，第二个参数是要从这个列表中抽样的个数。它的返回值也是一个列表，表示我们抽样的结果。

（2）shutil.move(source,dest)：source表示的源地址文件，dest表示的是目标地址文件

2.3 移动训练集的图片以及标签

# 将剩余的图片都移动到训练集中
def moveLeftFile(path_source, path_dest):
    for path_file in os.listdir(path_source):
        shutil.move(os.path.join(path_source,path_file),os.path.join(path_dest,path_file))
    print("移动完成")

# 移动训练图片和标签
moveLeftFile(path_source = path_images_dir, path_dest = path_datasets_train)
moveLeftFile(path_source = path_annotations_dir, path_dest = path_labels_train)