import os
import numpy as np
import tensorflow as tf


def get_train_data(train_dir, batch_size):
    train_images = np.load(os.path.join(train_dir, 'train_images.npy'))
    train_labels = np.load(os.path.join(train_dir, 'train_labels.npy'))
    print('train_images', train_images.shape, 'train_labels', train_labels.shape)

    dataset_train = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
    dataset_train = dataset_train.repeat().shuffle(10000).batch(batch_size)

    return dataset_train


def get_val_data(val_dir):
    test_images = np.load(os.path.join(val_dir, 'validation_images.npy'))
    test_labels = np.load(os.path.join(val_dir, 'validation_labels.npy'))
    print('validation_images', test_images.shape, 'validation_labels', test_labels.shape)

    dataset_test = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
    return dataset_test

[((11, 20, 11, 84), 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', ({(11, 55, 11, 83): '(train_images, train_labels)'}, {}), '((train_images, train_labels))', True, 'import tensorflow as tf\n'), ((22, 19, 22, 81), 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', ({(22, 54, 22, 80): '(test_images, test_labels)'}, {}), '((test_images, test_labels))', True, 'import tensorflow as tf\n'), ((7, 27, 7, 70), 'os.path.join', 'os.path.join', ({(7, 40, 7, 49): 'train_dir', (7, 51, 7, 69): '"""train_images.npy"""'}, {}), "(train_dir, 'train_images.npy')", False, 'import os\n'), ((8, 27, 8, 70), 'os.path.join', 'os.path.join', ({(8, 40, 8, 49): 'train_dir', (8, 51, 8, 69): '"""train_labels.npy"""'}, {}), "(train_dir, 'train_labels.npy')", False, 'import os\n'), ((18, 26, 18, 72), 'os.path.join', 'os.path.join', ({(18, 39, 18, 46): 'val_dir', (18, 48, 18, 71): '"""validation_images.npy"""'}, {}), "(val_dir, 'validation_images.npy')", False, 'import os\n'), ((19, 26, 19, 72), 'os.path.join', 'os.path.join', ({(19, 39, 19, 46): 'val_dir', (19, 48, 19, 71): '"""validation_labels.npy"""'}, {}), "(val_dir, 'validation_labels.npy')", False, 'import os\n')]