42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
import os
|
||
import random
|
||
import shutil
|
||
|
||
def split_dataset(input_folder, output_folder):
|
||
# 创建输出子目录
|
||
subdatasets = [os.path.join(output_folder, f'subdataset_{i+1}') for i in range(3)]
|
||
for subdataset in subdatasets:
|
||
os.makedirs(subdataset, exist_ok=True)
|
||
|
||
# 遍历每个类别文件夹
|
||
for root, dirs, _ in os.walk(input_folder):
|
||
for category in dirs:
|
||
category_folder = os.path.join(root, category)
|
||
images = [os.path.join(category_folder, f) for f in os.listdir(category_folder)
|
||
if f.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]
|
||
|
||
# 如果图像数量少于 1200，进行补充
|
||
if len(images) < 1200:
|
||
images.extend(random.choices(images, k=1200 - len(images)))
|
||
|
||
# 打乱图像顺序
|
||
random.shuffle(images)
|
||
|
||
# 为每个子数据集创建类别子文件夹
|
||
for subdataset_path in subdatasets:
|
||
category_subfolder = os.path.join(subdataset_path, category)
|
||
os.makedirs(category_subfolder, exist_ok=True)
|
||
|
||
# 将图像分配到子数据集中，每个子数据集 400 张
|
||
for i, image_path in enumerate(images):
|
||
subdataset_index = i // 400
|
||
subdataset_path = subdatasets[subdataset_index]
|
||
category_subfolder = os.path.join(subdataset_path, category)
|
||
shutil.copy(image_path, category_subfolder)
|
||
|
||
print(f'Dataset split into 3 subdatasets with 400 images per category at {output_folder}')
|
||
|
||
if __name__ == "__main__":
|
||
input_folder = 'L:/Grade_datasets/MOVE_BACKGROUND'
|
||
output_folder = 'L:/Grade_datasets/SPLIT'
|
||
split_dataset(input_folder, output_folder) |