数据集划分训练测试脚本

import os
import shutil
import random


curPath = os.getcwd()
fullImageDir = "./fireDetectionDATA/"

trainImageDir = "./train/"
valImageDir= "./test/"

if not os.path.exists(trainImageDir):
    os.mkdir(trainImageDir)
if not os.path.exists(valImageDir):
    os.mkdir(valImageDir)

imageList = os.listdir(fullImageDir)
image_num = len(imageList)
# shuffle all the data
random.shuffle(imageList)

count = 0
trainRatio = 0.8  # CHANGE
train_num = trainRatio * image_num

for i in imageList:
    count += 1
    if count <= train_num:
        shutil.copy(fullImageDir + i, trainImageDir)
    else:
        shutil.copy(fullImageDir + i, valImageDir)