# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# 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 get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/ctaugment.py
# OPS = {}
# OP = namedtuple('OP', ('f', 'bins'))
# def register(*bins):
#     def wrap(f):
# def apply(x, ops):
#     def __init__(self, depth=2, th=0.85, decay=0.99):
#     def rate_to_p(self, rate):
#     def policy(self, probe):
#     def update_rates(self, policy, proximity):
#     def stats(self):
# def _enhance(x, op, level):
# def _imageop(x, op, level):
# def _filter(x, op, level):
# def autocontrast(x, level):
# def blur(x, level):
# def brightness(x, brightness):
# def color(x, color):
# def contrast(x, contrast):
# def cutout(x, level):
# def equalize(x, level):
# def invert(x, level):
# def identity(x):
# def posterize(x, level):
# def rescale(x, scale, method):
# def rotate(x, angle):
# def sharpness(x, sharpness):
# def shear_x(x, shear):
# def shear_y(x, shear):
# def smooth(x, level):
# def solarize(x, th):
# def translate_x(x, delta):
# def translate_y(x, delta):
# class CTAugment:
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


import collections
import functools
import itertools
import multiprocessing
import random
import numpy as np
import tensorflow as tf
from absl import flags
from libml import utils, ctaugment
from libml.utils import EasyDict
from third_party.auto_augment import augmentations, policies

    def queue_images(self, batch):
        args = []
        image = batch['image']
        # Fixed magnitude
        policy = lambda: [(op, 1.0, self.magnitude) for op in np.random.choice(RANDOM_POLICY_OPS, self.nops)]
        if image.ndim == 4:
            for x in range(image.shape[0]):
                args.append((image[x:x + 1], [policy()], self.size))
        else:
            for x in image[:, 1:]:
                args.append((x, [policy() for _ in range(x.shape[0])], self.size))
        self.queue.append(PoolEntry(payload=POOL.imap(self.numpy_apply_policies, args), batch=batch))


class AugmentPoolCTA(AugmentPool):

    def __init__(self, get_samples):
        init_pool()
        self.get_samples = get_samples
        self.queue = []
        self.fill_queue()

    @staticmethod
    def numpy_apply_policies(arglist):
        x, cta, probe = arglist
        if x.ndim == 3:
            assert probe
            policy = cta.policy(probe=True)
            return dict(policy=policy,

# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# 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 get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/ctaugment.py
# OPS = {}
# OP = namedtuple('OP', ('f', 'bins'))
# def register(*bins):
#     def wrap(f):
# def apply(x, ops):
#     def __init__(self, depth=2, th=0.85, decay=0.99):
#     def rate_to_p(self, rate):
#     def policy(self, probe):
#     def update_rates(self, policy, proximity):
#     def stats(self):
# def _enhance(x, op, level):
# def _imageop(x, op, level):
# def _filter(x, op, level):
# def autocontrast(x, level):
# def blur(x, level):
# def brightness(x, brightness):
# def color(x, color):
# def contrast(x, contrast):
# def cutout(x, level):
# def equalize(x, level):
# def invert(x, level):
# def identity(x):
# def posterize(x, level):
# def rescale(x, scale, method):
# def rotate(x, angle):
# def sharpness(x, sharpness):
# def shear_x(x, shear):
# def shear_y(x, shear):
# def smooth(x, level):
# def solarize(x, th):
# def translate_x(x, delta):
# def translate_y(x, delta):
# class CTAugment:
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


# Path: libml/augment.py
import collections
import functools
import itertools
import multiprocessing
import random
import numpy as np
import tensorflow as tf
from absl import flags
from libml import utils, ctaugment
from libml.utils import EasyDict
from third_party.auto_augment import augmentations, policies


    def queue_images(self, batch):
        args = []
        image = batch['image']
        # Fixed magnitude
        policy = lambda: [(op, 1.0, self.magnitude) for op in np.random.choice(RANDOM_POLICY_OPS, self.nops)]
        if image.ndim == 4:
            for x in range(image.shape[0]):
                args.append((image[x:x + 1], [policy()], self.size))
        else:
            for x in image[:, 1:]:
                args.append((x, [policy() for _ in range(x.shape[0])], self.size))
        self.queue.append(PoolEntry(payload=POOL.imap(self.numpy_apply_policies, args), batch=batch))


class AugmentPoolCTA(AugmentPool):

    def __init__(self, get_samples):
        init_pool()
        self.get_samples = get_samples
        self.queue = []
        self.fill_queue()

    @staticmethod
    def numpy_apply_policies(arglist):
        x, cta, probe = arglist
        if x.ndim == 3:
            assert probe
            policy = cta.policy(probe=True)
            return dict(policy=policy,

                        probe=ctaugment.apply(x, policy),