Python 取得した画像を学習用とテスト用に分けよう!

#画像に処理をするためのライブラリ
from PIL import Image

#ディレクトリやファイルに処理をするためのライブラリ
import os, glob

import numpy as np

#データを学習用とテスト用に分けるためのライブラリ
from sklearn import model_selection

#クラスを作成
veg_classes = ["cabbage", "lettuce"]


#クラスの数（＝2）
num_veg_classes = len(veg_classes)


#画像のサイズを指定
image_size = 50

#画像データの配列を用意
X = []

#ラベルのデータを用意
Y = []

#画像データとラベルデータに分割
for index, classlabel in enumerate(veg_classes):
    
    #画像のディレクトリ名
    photo_dir = "./" + classlabel
    
    #画像を一枚ずつ取得
    files = glob.glob(photo_dir + "/*.jpg")
    
    #写真をRGB化し，変更を加える
    for i, file in enumerate(files):
        
        #もし，150枚以上なら，終了
        if i>=150: break
    
        #画像を開く
        image = Image.open(file)
        
        #RGB化する
        image = image.convert("RGB")
        
        #サイズを50×50にリサイズする
        image = image.resize((image_size, image_size))
        
        #画像データを数字に変換する
        data = np.asarray(image)
        
        X.append(data)

        Y.append(index)

#XとYをnumpy形式に変換する（TensorFlowでは，numpy形式の方が扱いやすい）
X = np.array(X)

Y = np.array(Y)

#学習用とテスト用に分割
X_train, X_test, Y_train, Y_test = model_selection.train_test_split(X, Y)

#学習用とテスト用のデータをまとめておく
xy = (X_train, X_test, Y_train, Y_test)

#Numpy形式でファイルを保存しておく
np.save("./vegetables.npy", xy)