CVer肯定都遇到过要划分数据集的问题,我在这里自己写了一个,只要你有一个所有图片的文件夹,然后将路径替换成自己的即可:
# 将数据集随机划分为训练集和验证集,测试集
import os
import random
import shutil
from tqdm import tqdm
image_path = r'E:\DCI(first)_3\DCI_Crop2\image\\' # 源图片文件夹路径
mask_path = r'E:\DCI(first)_3\DCI_Crop2\mask\\' # 标签文件夹路径
train_images = r'E:\DCI(first)_3\DCI_Split2\train\images' # 划分后训练集图片的保存路径
train_labels = r'E:\DCI(first)_3\DCI_Split2\train\labels'
val_images = r'E:\DCI(first)_3\DCI_Split2\val\images'
val_labels = r'E:\DCI(first)_3\DCI_Split2\val\labels'
test_images = r'E:\DCI(first)_3\DCI_Split2\test\images'
test_labels = r'E:\DCI(first)_3\DCI_Split2\test\labels'
if not os.path.exists(train_images):
os.makedirs(train_images)
if not os.path.exists(train_labels):
os.makedirs(train_labels)
if not os.path.exists(val_images):
os.makedirs(val_images)
if not os.path.exists(val_labels):
os.makedirs(val_labels)
if not os.path.exists(test_images):
os.makedirs(test_images)
if not os.path.exists(test_labels):
os.makedirs(test_labels)
train_rate = 0.7 # 自定义抽取图片的比例,比方说100张抽10张,那就是0.1
val_rate = 0.15
test_rate = 0.15
# 求训练集
pathDir = os.listdir(image_path) # 取图片的原始路径
print('数据集总共有图片:', len(pathDir))
print(
'划分比例如下:训练集:{},验证集:{},测试集:{}'.format(int(len(pathDir) * train_rate), int(len(pathDir) * val_rate),
int(len(pathDir) * test_rate)))
picknumber = int(len(pathDir) * train_rate) # 按照rate比例从文件夹中取一定数量图片
train_sample = random.sample(pathDir, picknumber) # 随机选取picknumber数量的样本图片
print('训练集的大小为:', len(train_sample))
# 复制为训练集
for name in tqdm(train_sample):
shutil.copy(image_path + name, train_images + "\\" + name)
shutil.copy(mask_path + name, train_labels + "\\" + name)
# 求出原数据集不含训练集
all_images = os.listdir(image_path)
remaining_image = []
for file in all_images:
if file not in train_sample:
remaining_image.append(file)
# 求验证集
picknumber2 = int(len(remaining_image) * val_rate / (val_rate + test_rate)) # 按照rate比例从文件夹中取一定数量图片
val_sample = random.sample(remaining_image, picknumber2) # 随机选取picknumber数量的样本图片
print('验证集的大小为:', len(val_sample))
# 复制为验证集
for file in tqdm(val_sample):
shutil.copy(image_path + file, val_images + "\\" + file)
shutil.copy(mask_path + file, val_labels + "\\" + file)
test_sample = []
for file in remaining_image:
if file not in val_sample:
test_sample.append(file)
print('测试集的大小为:', len(test_sample))
# 复制为测试集
for file in tqdm(test_sample):
shutil.copy(image_path + file, test_images + "\\" + file)
shutil.copy(mask_path + file, test_labels + "\\" + file)