import os
import cv2
import random
import numpy as np
from tensorflow.keras.utils import to_categorical
from scripts.consts import class_dict


def get_data(path, split=0.2):

    X, y = [], []

    for directory in os.listdir(path):

        dirpath = os.path.join(path, directory)
        print(directory, len(os.listdir(dirpath)))

        for file in os.listdir(dirpath):

            filepath = os.path.join(dirpath, file)
            img = cv2.imread(filepath, cv2.IMREAD_UNCHANGED)
            
            if img.shape != (360, 363, 3):
                img = cv2.resize(img, (360, 363), cv2.INTER_CUBIC)

            X.append(img)
            y.append(class_dict[directory])

    data = list(zip(X, y))

    random.shuffle(data)

    X, y = zip(*data)

    num_train = int((1.0 - split) * len(y))

    X_train, X_valid = np.array(X[:num_train]).astype(
        'float32'), np.array(X[num_train:]).astype('float32')
    y_train, y_valid = np.array(
        y[:num_train]).reshape(-1, 1), np.array(y[num_train:]).reshape((-1, 1))

    X_train = X_train / 255.0
    X_valid = X_valid / 255.0

    y_train, y_valid = to_categorical(y_train), to_categorical(y_valid)

    print(X_train.shape, y_train.shape)
    print(X_valid.shape, y_valid.shape)

    return X_train, y_train, X_valid, y_valid

[((13, 21, 13, 37), 'os.listdir', 'os.listdir', ({(13, 32, 13, 36): 'path'}, {}), '(path)', False, 'import os\n'), ((31, 4, 31, 24), 'random.shuffle', 'random.shuffle', ({(31, 19, 31, 23): 'data'}, {}), '(data)', False, 'import random\n'), ((15, 18, 15, 47), 'os.path.join', 'os.path.join', ({(15, 31, 15, 35): 'path', (15, 37, 15, 46): 'directory'}, {}), '(path, directory)', False, 'import os\n'), ((18, 20, 18, 39), 'os.listdir', 'os.listdir', ({(18, 31, 18, 38): 'dirpath'}, {}), '(dirpath)', False, 'import os\n'), ((45, 23, 45, 46), 'tensorflow.keras.utils.to_categorical', 'to_categorical', ({(45, 38, 45, 45): 'y_train'}, {}), '(y_train)', False, 'from tensorflow.keras.utils import to_categorical\n'), ((45, 48, 45, 71), 'tensorflow.keras.utils.to_categorical', 'to_categorical', ({(45, 63, 45, 70): 'y_valid'}, {}), '(y_valid)', False, 'from tensorflow.keras.utils import to_categorical\n'), ((20, 23, 20, 50), 'os.path.join', 'os.path.join', ({(20, 36, 20, 43): 'dirpath', (20, 45, 20, 49): 'file'}, {}), '(dirpath, file)', False, 'import os\n'), ((21, 18, 21, 60), 'cv2.imread', 'cv2.imread', ({(21, 29, 21, 37): 'filepath', (21, 39, 21, 59): 'cv2.IMREAD_UNCHANGED'}, {}), '(filepath, cv2.IMREAD_UNCHANGED)', False, 'import cv2\n'), ((16, 29, 16, 48), 'os.listdir', 'os.listdir', ({(16, 40, 16, 47): 'dirpath'}, {}), '(dirpath)', False, 'import os\n'), ((24, 22, 24, 66), 'cv2.resize', 'cv2.resize', ({(24, 33, 24, 36): 'img', (24, 38, 24, 48): '(360, 363)', (24, 50, 24, 65): 'cv2.INTER_CUBIC'}, {}), '(img, (360, 363), cv2.INTER_CUBIC)', False, 'import cv2\n'), ((37, 23, 37, 46), 'numpy.array', 'np.array', ({(37, 32, 37, 45): 'X[:num_train]'}, {}), '(X[:num_train])', True, 'import numpy as np\n'), ((38, 20, 38, 43), 'numpy.array', 'np.array', ({(38, 29, 38, 42): 'X[num_train:]'}, {}), '(X[num_train:])', True, 'import numpy as np\n'), ((39, 23, 40, 22), 'numpy.array', 'np.array', ({(40, 8, 40, 21): 'y[:num_train]'}, {}), '(y[:num_train])', True, 'import numpy as np\n'), ((40, 39, 40, 62), 'numpy.array', 'np.array', ({(40, 48, 40, 61): 'y[num_train:]'}, {}), '(y[num_train:])', True, 'import numpy as np\n')]