【数据集】|数据集拆分

import os
import shutil
import random
 
 
root_dir = "/home/code/all_in_one/"
imgs_dir = os.path.join(root_dir, "images/")
labs_dir = os.path.join(root_dir, "labels/")
img_train = "/home/code/all_train/images"
img_val = "/home/code/all_val/images"
lab_train = "/home/code/all_train/labels"
lab_val = "/home/code/all_val/labels"
 
imgs_list = os.listdir(imgs_dir)
labs_list = os.listdir(labs_dir)
imgs_list.sort()
#print(imgs_list)
labs_list.sort()
lens = len(imgs_list)
idx_list = [i for i in range(lens)]
random.shuffle(idx_list)  # 将列表进行了打乱
 
for i, idx in enumerate(idx_list):
    img_name = imgs_list[idx]
    lab_name = labs_list[idx]
    img_path = os.path.join(imgs_dir, img_name)
    lab_path = os.path.join(labs_dir, lab_name)
    print(img_path)
    print(lab_path)
    if i < lens * 0.98:  # 作为训练集
        img_train_path = os.path.join(img_train, img_name)
        lab_train_path = os.path.join(lab_train, lab_name)
        shutil.copyfile(img_path, img_train_path)
        shutil.copyfile(lab_path, lab_train_path)
    else:
        img_val_path = os.path.join(img_val, img_name)
        lab_val_path = os.path.join(lab_val, lab_name)
        shutil.copyfile(img_path, img_val_path)
        shutil.copyfile(lab_path, lab_val_path)