import os 
import cv2
import numpy as np

#make the function to load all the images in the dataset

def loadDataset(path):
  for _,dirs,_ in os.walk(path):
    X=[];
    y=[];
    for dir in dirs:
        for _,_,files in os.walk(path+dir+"/"):
            for file in files:
                if file.find('.png')<0:
                       continue     
              #  print("Loading:"+path+dir+"/"+file)
                img = cv2.imread(path+dir+"/"+file)
                resized = cv2.resize(img, (100,100), interpolation = cv2.INTER_AREA) #resize image 
                gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY) #change to grayscale
                X.append(gray)
                y.append(dir)
    return X,y


X_train, y_train=loadDataset('/home/mostafa/Desktop/ML- Labs/CHALLENGE_2/dataset/train/')
X_train = np.array(X_train)
y_train = np.array(y_train)
X_train = X_train.astype('float') / 255

X,y = loadDataset('/home/mostafa/Desktop/ML- Labs/CHALLENGE_2/dataset/validationandtest/')
X=np.array(X)
y = np.array(y) 
X = X.astype('float') / 255

from sklearn.model_selection import train_test_split
X_test, X_val, y_test, y_val = train_test_split(X,y, test_size=0.2)

print('X_train shape: ',X_train.shape,y_train.shape)
print('X_val shape: ',X_val.shape, y_val.shape)
print('X_test shape: ',X_test.shape, y_test.shape)