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

[((345, 409), 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', (['(train_images, train_labels)'], {}), '((train_images, train_labels))\n', (379, 409), True, 'import tensorflow as tf\n'), ((799, 861), 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', (['(test_images, test_labels)'], {}), '((test_images, test_labels))\n', (833, 861), True, 'import tensorflow as tf\n'), ((125, 168), 'os.path.join', 'os.path.join', (['train_dir', '"""train_images.npy"""'], {}), "(train_dir, 'train_images.npy')\n", (137, 168), False, 'import os\n'), ((197, 240), 'os.path.join', 'os.path.join', (['train_dir', '"""train_labels.npy"""'], {}), "(train_dir, 'train_labels.npy')\n", (209, 240), False, 'import os\n'), ((567, 613), 'os.path.join', 'os.path.join', (['val_dir', '"""validation_images.npy"""'], {}), "(val_dir, 'validation_images.npy')\n", (579, 613), False, 'import os\n'), ((641, 687), 'os.path.join', 'os.path.join', (['val_dir', '"""validation_labels.npy"""'], {}), "(val_dir, 'validation_labels.npy')\n", (653, 687), False, 'import os\n')]