python 将数据集随机抽取并改名保存

import os
import numpy as np
import random
if __name__ == '__main__':
    random.seed(200)
    index = np.arange(0,7620)
    random.shuffle(index)
    print(index)
    train_num = 6096
    test_num = 762
    val_num = 762
    train_index = index[0:train_num]
    test_index = index[train_num:train_num+test_num]
    val_index = index[train_num+test_num:train_num+test_num+val_num]
    # print(len(train_index),len(test_index),len(val_index))
    path = r"F:\pycharm_doc\tools\whu_out"
    filelst = os.listdir(path)
    print(filelst)
    # ['after', 'before', 'change_label']
    for name in filelst:
        # ['1.jpg', '10.jpg', '100.jpg', '1000.jpg', '1001.jpg',...
        imaglst = os.listdir(os.path.join(path,name))
        # train_imaglst = imaglst[train_index]
        for i in train_index:
            train_name = imaglst[i]
            old = os.path.join(path, name, train_name)
            new = os.path.join(path, name, 'train_'+train_name)
            print(old, '->', new)
            os.rename(old, new)
        # tes_imaglst = imaglst[test_index]
        for i in test_index:
            test_name = imaglst[i]
            old = os.path.join(path, name, test_name)
            new = os.path.join(path, name, 'test_'+test_name)
            print(old, '->', new)
            os.rename(old, new)
        for i in val_index:
            val_name = imaglst[i]
            old = os.path.join(path, name, val_name)
            new = os.path.join(path, name, 'val_'+val_name)
            print(old, '->', new)
            os.rename(old, new)