將不同類的圖像分割成Python中的訓練和測試集

Question

我是python的新手，但我想在python中實現深度學習工具。我收集了一組不同類別或類別的圖像。我的工作是使用卷積網絡執行圖像分類。第一步是將這些圖像分成兩組進行訓練和測試。然後，我將加載這些圖像並進行一些預處理，然後將它們饋送到網絡中。我現在只對第一部分感到好奇。

我後我的工作是：

import numpy as np 
from os import listdir 

def getPaths(rootPath): 
    folderList = listdir(rootPath) 
    imgPaths = [] 
    categories = [] 
    for folders in folderList: 
    imgPath = os.path.join(rootPath,folders) 
    imgPaths.append(imgPath) 
    categories.append(folders) 
    return (imgPaths, categories) 

def getImgPaths(rootPath, p):  
temp = getPaths(rootPath) 
folderPaths, categories = temp 

trainImgPaths = [] 
trainLabels = [] 
testImgPaths = [] 
testLabels = [] 
for ii in range(len(folderPaths)): 
    temp2 = getPaths(folderPaths[ii]) 
    imgPaths = temp2[0] 
    randIdx = np.random.permutation(len(imgPaths)) 
    trainIdx = randIdx[:int(p*len(imgPaths))] 
    testIdx = [idx for idx in randIdx if not idx in trainIdx] 
    trainPaths = [imgPaths[kk] for kk in trainIdx] 
    testPaths = [imgPaths[kk] for kk in testIdx] 
    trainCat = [categories[ii] for jj in xrange(len(trainPaths))] 
    testCat = [categories[ii] for jj in xrange(len(testPaths))] 

    trainImgPaths.extend(trainPaths) 
    testImgPaths.extend(testPaths) 
    trainLabels.extend(trainCat) 
    testLabels.extend(testCat) 
    return (trainImgPaths, trainLabels, testImgPaths, testLabels)  

的代碼可以工作，但似乎有點麻煩。

Answer 1

import os, random 

def getImagePaths(imgroot, cats, pot): # please excuse the naming 
    trainImagePaths = [] 
    testImagePaths = [] 
    trainlabels = [] 
    testlabels = [] 
    for cat in cats: 
     files = os.listdir(os.path.join(imgroot, cat)) 
     split = int(pot*len(files)) 
     trainImagePaths.extend(files[:split]) 
     testImagePths.extend(files[split:] 
     trainlabels.extend([cat]*split) 
     testlabels.extend([cat]*len(files)-split) 

    # optionally, shuffle 

    return trainImagePaths, testImagePaths, trainlabels, testLabels