随机划分数据集

import os  
import random  
import math  
import shutil  
  
def data_split(old_path, new_path):  
    if os.path.exists(new_path) == 0:  
        os.makedirs(new_path)  
    for dirpath, dirnames, filenames in os.walk(old_path):  
        for filename in filenames:  
            if filename.endswith('.jpg'):  
                filepath = os.path.join(dirpath, filename)  
                random.shuffle(filenames)  
                n = len(filenames)  
                train_ratio = 0.8  
                val_ratio = 0.1  
                test_ratio = 0.1  
                train_count = int(n * train_ratio)  
                val_count = int(n * val_ratio)  
                test_count = n - train_count - val_count  
                train_subdir = os.path.join(new_path, 'train_set')  
                val_subdir = os.path.join(new_path, 'val_set')  
                test_subdir = os.path.join(new_path, 'test_set')  
                if not os.path.exists(train_subdir):  
                    os.makedirs(train_subdir)  
                if not os.path.exists(val_subdir):  
                    os.makedirs(val_subdir)  
                if not os.path.exists(test_subdir):  
                    os.makedirs(test_subdir)  
                train_files = random.sample(filenames, train_count)  
                val_files = random.sample(filenames, val_count)  
                test_files = filenames[-test_count:]  
                for i, train_file in enumerate(train_files):  
                    shutil.copy2(filepath, os.path.join(train_subdir, train_file))  
                for i, val_file in enumerate(val_files):  
                    shutil.copy2(filepath, os.path.join(val_subdir, val_file))  
                for i, test_file in enumerate(test_files):  
                    shutil.copy2(filepath, os.path.join(test_subdir, test_file))  
  
if __name__ == '__main__':  
    data_path = '/home/uto/AIV_5_4筛/'  
    data_split(data_path, 'data')