python将一个文件夹下图片和对应的json标签按比例分为训练集,验证集和测试集三个文件夹下

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# 将一个文件夹下图片按比例分在三个文件夹下
import os
import random
import shutil
from shutil import copy2

datadir_normal = "./image/"

all_data = os.listdir(datadir_normal)  # (图片文件夹)
all_data_img = []
for i in all_data:
    if i.endswith(".jpeg"):
        all_data_img.append(i)
num_all_data = len(all_data_img)
print("num_all_data: " + str(num_all_data))
index_list = list(range(num_all_data))
print(index_list)
random.shuffle(index_list)
num = 0

trainDir = "./train/"  # (将训练集放在这个文件夹下)
if not os.path.exists(trainDir):
    os.mkdir(trainDir)

validDir = './val/'  # (将验证集放在这个文件夹下)
if not os.path.exists(validDir):
    os.mkdir(validDir)

testDir = './test/'  # (将测试集放在这个文件夹下)
if not os.path.exists(testDir):
    os.mkdir(testDir)

for i in index_list:
    fileName = os.path.join(datadir_normal, all_data_img[i])
    json_file = all_data_img[i].split('.')[0] + ".json"
    json_fileName = os.path.join(datadir_normal,json_file)
    if num < num_all_data * 0.8:
        copy2(fileName, trainDir)
        copy2(json_fileName,trainDir)
    elif num > num_all_data * 0.8 and num < num_all_data * 1:
        # print(str(fileName))
        copy2(fileName, validDir)
        copy2(json_fileName,validDir)
    else:
        copy2(fileName, testDir)
        copy2(json_fileName,testDir)
    num += 1

你可能感兴趣的:(python,python,ubuntu,linux)