随机分配训练集测试集和验证集

这个代码更加适用于,一张图片对应一个标注文件的数据集。通过随机分配之后,将全部文件分为train、test、val以及对应的标注数据也会分开。

如果您只有图像没有标注数据,那就更容易实现。本文就不再赘述。

下面是代码,核心部分就是使用了一次随机排序。

import os  # os是用来切换路径和创建文件夹的。
import random
from shutil import copy  # shutil 是用来复制黏贴文件的

# "blossom_end_rot", "graymold", "powdery_mildew", "spider_mite", "spotting_disease"

img_path = "/Users/Dong/Desktop/dataPath/data"
labels_path = "/Users/Dong/Desktop/dataPath/labels"

filelist_img_name = []
filelist_label_name = []

for root, dir, filenames in os.walk(img_path):
    for filename in filenames:
        if '.jpg' in filename:
            file_path = os.path.join(root, filename)
            print(file_path)
            filelist_img_name.append(file_path)
        else:
            pass
            # print("it is not a jpg file.", filename)

# .jpg.mat是我的标注文件的后缀名,您修改为自己的xml或者json或者其他类型即可。
for root, dir, filenames in os.walk(labels_path):
    for filename in filenames:
        if '.jpg.mat' in filename:
            filelist_label_name.append(filename.split('.jpg.mat')[0])
        else:
            pass
            # print("it is not a label file.", filename)

print(len(filelist_img_name))
print(len(filelist_label_name))
# 我一共含有580张图片,因此 348 ,464,464-580这几个数字是我自己计算的。我的比例是60% 20% 20%
# 如果您有10000张图片和10000张标注,那么您可以设置为6000,8000,8000:,以此类推
# 我没有写自动创建文件夹的那部分代码,因此你要自己首先把相关文件夹创建好。
random.shuffle(filelist_img_name)
filelist_img_name = filelist_img_name
for tr in filelist_img_name[0:348]:
    copy(tr, '/Users/Dong/Desktop/dataPath_new/data/train/')
    tr = tr.replace('data/', 'labels/')
    tr = tr.replace('.jpg', '.jpg.mat')
    copy(tr, '/Users/Dong/Desktop/dataPath_new/labels/train/')

for val in filelist_img_name[348:464]:
    copy(val, '/Users/Dong/Desktop/dataPath_new/data/val/')
    val = val.replace('data/', 'labels/')
    val = val.replace('.jpg', '.jpg.mat')
    copy(val, '/Users/Dong/Desktop/dataPath_new/labels/val/')

for test in filelist_img_name[464:]:
    copy(test, '/Users/Dong/Desktop/dataPath_new/data/test/')
    test = test.replace('data/', 'labels/')
    test = test.replace('.jpg', '.jpg.mat')
    copy(test, '/Users/Dong/Desktop/dataPath_new/labels/test/')

猜你喜欢

转载自blog.csdn.net/Dongjiuqing/article/details/121647592
今日推荐