[python]将图片按比例分在多个文件夹(7:3分成训练集和测试集)

先看看原数据集:
第一层

第二层

其他类似~

再看看修改后数据集:
第一层:

第二层:

第三层:

test类似,这里不做展示-_-

然后上代码

# -*- coding: utf-8 -*-
"""
Created on Mon May 25 23:31:16 2020

@author: 吴
"""
#代码没什么地方难理解的,没看懂建议多print参数看看
import os
import random
import shutil
from shutil import copy2


category = ['可回收垃圾', '有害垃圾', '厨余垃圾', '其他垃圾']
#要保存的根目录
saveTrainPATH = r'D:\百步梯\图像识别\garbageClassifier\IMAGEDATA\train'
saveTestPATH = r'D:\百步梯\图像识别\garbageClassifier\IMAGEDATA\test'

if __name__=="__main__":
    
    for j in range(4):
        PATH = r'D:\百步梯\图像识别\垃圾目录' + '/'+  category[j]
    
        #子文件夹
        for childPATH in os.listdir(PATH):
            #子文件夹路径
            NewChildPATH = PATH + '/'+ str(childPATH)
            print(NewChildPATH)
            trainfiles = os.listdir(NewChildPATH)
            num_image = len(trainfiles)
            
            # print( NewChildPATH + "   \t num_image: " + str(num_image) )
            
            index_list = list(range(num_image))
            # print(index_list)
            # random.shuffle(index_list)
            num = 0
            
            #保存trian的路径-----------------------------
            trainDir = saveTrainPATH + '/' + category[j]
            #先判断是否存在这个文件夹
            if not os.path.exists(trainDir):                   
                os.mkdir(str(trainDir))
                
            
            childTrainDir = trainDir + '/' + str(childPATH)
            #判断子子文件夹是否存在,若不存在则创建(老套娃了~~)
            if not os.path.exists(childTrainDir):
                os.mkdir(str(childTrainDir))
                
            
            #保存test的路径---------------------------------
            testDir = saveTestPATH + '/' + '/' + category[j]   
            if not os.path.exists(testDir):
                os.mkdir(str(testDir))
                
                
            childTestDir = testDir + '/' + str(childPATH)
            #判断子子文件夹是否存在,若不存在则创建
            if not os.path.exists(childTestDir):
                os.mkdir(str(childTestDir))
                
            for i in index_list:
                fileName = os.path.join(NewChildPATH, trainfiles[i])
                if num < num_image*0.7:
                    copy2(fileName, childTrainDir)  #复制过去,不改变原来目录的图片
                else:
                    copy2(fileName, childTestDir)
                num += 1
               
            
            print(trainDir,'\n',testDir )

至此数据集处理终于完成了,有关数据集的处理可以看看我之前写的博客

1.python批量修改图片格式和名称
2.python文件夹下多个子文件夹,对图片批量修改格式和名称
3.python批量修改图片尺寸(含多个文件夹)