from PIL import Image
import os, glob
import numpy as np
from sklearn import model_selection
classes = ["shs","bns","amatsuji","screw","skyscraper","column"]
num_classes = len(classes)
image_size = 50
num_testdata = 100

#画像読み込み
X_train = []
X_test = []
Y_train = []
Y_test = []

for index, classlabel in enumerate(classes):
#X [0]に値が入っていない場合はdirの指示先が外れている。
photos_dir = "d:/product_ai/" + classlabel
files = glob.glob(photos_dir + "/*.jpg")
for i, file in enumerate(files):
if i >= 200: break
image = Image.open(file)
image = image.convert("RGB")
image = image.resize((image_size, image_size))
data = np.asarray(image)

if i < num_testdata:
X_test.append(data)
Y_test.append(index)
else:
#X_train .append(data)
#Y_train .append(index)

for angle in range(-20,20,5):
#回転
img_r = image.rotate(angle)
data = np.asarray(img_r)
X_train.append(data)
Y_train.append(index)
#反転
img_trans = img_r.transpose(Image.FLIP_LEFT_RIGHT)
data = np.asarray(img_trans)
X_train.append(data)
Y_train.append(index)

# data = np.asarray(image)
# X.append(data)
# Y.append(index)

X_train = np.array(X_train)
X_test = np.array(X_test)
y_train = np.array(Y_train)
y_test = np.array(Y_test)
#==================================================
#-------------------train, testのSplit
#X_train , X_test, y_train, y_test = model_selection.train_test_split(X,Y)
xy = (X_train, X_test, y_train, y_test)
np.save("d:/product_ai/product_aug.npy", xy)