随机划分训练集和测试集python代码

"""
本程序是随机划分训练集和测试集的
"""
import os
import random
import shutil


def howmuch(path):
    sum =0
    for i in os.listdir(path):
        sum+=1
    return sum

def move(src,dst,number):
    #随机生成要移动的数字列表
    list = random.sample(range(0, sum),number)

    for i in list :
         # 图像后缀（自己找图片的名称）
        img = 'or'+str(i)+".jpg"
        txt = 'or'+str(i)+".txt"
        #移动
        shutil.move(os.path.join(src[0],img),os.path.join(dst[0],img))
        shutil.move(os.path.join(src[1], txt),os.path.join(dst[1], txt))

if __name__ == '__main__':


    Img_path = r'E:\Desktop\x\data\images' #图像路径
    Label_path = r'E:\Desktop\x\data\labels' #标签位置
    #数据集总数
    sum = howmuch(Img_path)
s
    #设置划分比例
    test = 0.2
    train = 0.8
    #移动文件
    #保存的路径
    Save_train = r'E:\Desktop\x\data\images\train' #图像路径
    Save_train_labels = r'E:\Desktop\x\data\images\train_labels'  # 图像路径
    Save_test = r'E:\Desktop\x\data\test\img'  # 图像路径
    Save_test_labels = r'E:\Desktop\x\data\test\labels'  # 图像路径
    #随机移动测试集
    #src是原来路径
    src =[Img_path,Label_path]
    #dst 移动后的路径
    dst =[Save_test,Save_test_labels]
    #number 为移动的个数
    number = int( test*sum)
    move(src,dst,number)