# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):


import glob
import itertools
import os
import numpy as np
import tensorflow as tf
from absl import flags
from tqdm import tqdm
from libml import utils
    def tf_get(index):
        def get(index):
            return images[index], labels[index]

        image, label = tf.py_func(get, [index], [tf.float32, tf.int64])
        return dict(image=image, label=label)

    dataset = tf.data.Dataset.range(len(data)).repeat()
    dataset = dataset.shuffle(len(data) if len(data) < FLAGS.shuffle else FLAGS.shuffle)
    return dataset.map(tf_get)


def augment_mirror(x):
    return tf.image.random_flip_left_right(x)


def augment_shift(x, w):
    y = tf.pad(x, [[w] * 2, [w] * 2, [0] * 2], mode='REFLECT')
    return tf.random_crop(y, tf.shape(x))


def augment_noise(x, std):
    return x + std * tf.random_normal(tf.shape(x), dtype=x.dtype)


def compute_mean_std(data: tf.data.Dataset):
    data = data.map(lambda x: x['image']).batch(1024).prefetch(1)
    data = data.make_one_shot_iterator().get_next()
    count = 0
    stats = []

# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):


# Path: libml/data.py
import glob
import itertools
import os
import numpy as np
import tensorflow as tf
from absl import flags
from tqdm import tqdm
from libml import utils

    def tf_get(index):
        def get(index):
            return images[index], labels[index]

        image, label = tf.py_func(get, [index], [tf.float32, tf.int64])
        return dict(image=image, label=label)

    dataset = tf.data.Dataset.range(len(data)).repeat()
    dataset = dataset.shuffle(len(data) if len(data) < FLAGS.shuffle else FLAGS.shuffle)
    return dataset.map(tf_get)


def augment_mirror(x):
    return tf.image.random_flip_left_right(x)


def augment_shift(x, w):
    y = tf.pad(x, [[w] * 2, [w] * 2, [0] * 2], mode='REFLECT')
    return tf.random_crop(y, tf.shape(x))


def augment_noise(x, std):
    return x + std * tf.random_normal(tf.shape(x), dtype=x.dtype)


def compute_mean_std(data: tf.data.Dataset):
    data = data.map(lambda x: x['image']).batch(1024).prefetch(1)
    data = data.make_one_shot_iterator().get_next()
    count = 0
    stats = []

    with tf.Session(config=utils.get_config()) as sess: