os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ['TF_GPU_ALLOCATOR'] = "cuda_malloc_async"
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)

def get_label(file_path):
    parts = tf.strings.split(file_path, os.path.sep)
    one_hot = parts[-2] == class_names
    return tf.argmax(one_hot)

def decode_img(img):
    img = tf.io.decode_image(img, channels=3, expand_animations = False)
    img = tf.image.resize(img, [244, 244])
    img = tf.cast(img, tf.float32)
    return img

def process_path(file_path):
    label = get_label(file_path)
    img = tf.io.read_file(file_path)
    img = decode_img(img)
    return img, label

def configure_for_performance(ds):
    ds = ds.batch(128)
    ds = ds.prefetch(buffer_size=tf.data.AUTOTUNE)
    return ds

files = tf.data.Dataset.list_files((data_dir + '/*/*.png'), shuffle=False)
files = files.shuffle(image_count, reshuffle_each_iteration=False)

val_size = int(image_count * 0.2)

train_files = files.skip(val_size)
val_files = files.take(val_size)

train_ds = train_files.interleave(lambda x: tf.data.Dataset.from_tensor_slices([x]), cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE)
train_ds = train_ds.map(process_path, num_parallel_calls=tf.data.AUTOTUNE)

val_ds = val_files.interleave(lambda x: tf.data.Dataset.from_tensor_slices([x]), cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE)
val_ds = val_ds.map(process_path, num_parallel_calls=tf.data.AUTOTUNE)

train_ds = configure_for_performance(train_ds)
val_ds = configure_for_performance(val_ds)