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

[((216, 232), 'os.listdir', 'os.listdir', (['path'], {}), '(path)\n', (226, 232), False, 'import os\n'), ((716, 736), 'random.shuffle', 'random.shuffle', (['data'], {}), '(data)\n', (730, 736), False, 'import random\n'), ((253, 282), 'os.path.join', 'os.path.join', (['path', 'directory'], {}), '(path, directory)\n', (265, 282), False, 'import os\n'), ((355, 374), 'os.listdir', 'os.listdir', (['dirpath'], {}), '(dirpath)\n', (365, 374), False, 'import os\n'), ((1121, 1144), 'tensorflow.keras.utils.to_categorical', 'to_categorical', (['y_train'], {}), '(y_train)\n', (1135, 1144), False, 'from tensorflow.keras.utils import to_categorical\n'), ((1146, 1169), 'tensorflow.keras.utils.to_categorical', 'to_categorical', (['y_valid'], {}), '(y_valid)\n', (1160, 1169), False, 'from tensorflow.keras.utils import to_categorical\n'), ((400, 427), 'os.path.join', 'os.path.join', (['dirpath', 'file'], {}), '(dirpath, file)\n', (412, 427), False, 'import os\n'), ((446, 488), 'cv2.imread', 'cv2.imread', (['filepath', 'cv2.IMREAD_UNCHANGED'], {}), '(filepath, cv2.IMREAD_UNCHANGED)\n', (456, 488), False, 'import cv2\n'), ((312, 331), 'os.listdir', 'os.listdir', (['dirpath'], {}), '(dirpath)\n', (322, 331), False, 'import os\n'), ((567, 611), 'cv2.resize', 'cv2.resize', (['img', '(360, 363)', 'cv2.INTER_CUBIC'], {}), '(img, (360, 363), cv2.INTER_CUBIC)\n', (577, 611), False, 'import cv2\n'), ((829, 852), 'numpy.array', 'np.array', (['X[:num_train]'], {}), '(X[:num_train])\n', (837, 852), True, 'import numpy as np\n'), ((881, 904), 'numpy.array', 'np.array', (['X[num_train:]'], {}), '(X[num_train:])\n', (889, 904), True, 'import numpy as np\n'), ((946, 969), 'numpy.array', 'np.array', (['y[:num_train]'], {}), '(y[:num_train])\n', (954, 969), True, 'import numpy as np\n'), ((995, 1018), 'numpy.array', 'np.array', (['y[num_train:]'], {}), '(y[num_train:])\n', (1003, 1018), True, 'import numpy as np\n')]