augment.py

# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Augmentations for images.
"""
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

FLAGS = flags.FLAGS
POOL = None
POLICIES = EasyDict(cifar10=policies.cifar10_policies(),
                    svhn=policies.svhn_policies(),
                    svhn_noextra=policies.svhn_policies())

RANDOM_POLICY_OPS = (
    'Identity', 'AutoContrast', 'Equalize', 'Rotate',
    'Solarize', 'Color', 'Contrast', 'Brightness',
    'Sharpness', 'ShearX', 'TranslateX', 'TranslateY',
    'Posterize', 'ShearY'
)
AUGMENT_ENUM = 'd x m aa'.split() + ['r%d_%d_%d' % (nops, mag, cutout) for nops, mag, cutout in
                                     itertools.product(range(1, 5), range(1, 16), range(0, 100, 25))]

flags.DEFINE_integer('K', 1, 'Number of strong augmentation for unlabeled data.')
flags.DEFINE_enum('augment', 'd.d',
                  [x + '.' + y for x, y in itertools.product(AUGMENT_ENUM, AUGMENT_ENUM)] +
                  [x + '.' + y + '.' + z for x, y, z in itertools.product(AUGMENT_ENUM, 'dx', AUGMENT_ENUM)],
                  'Dataset augmentation method (x=identity, m=mirror, d=default, aa=auto-augment, '
                  'rxyy=random augment with x ops and magnitude yy),'
                  'first is for labeled data, others are for unlabeled.')


def init_pool():
    global POOL
    if POOL is None:
        para = max(1, len(utils.get_available_gpus())) * FLAGS.para_augment
        POOL = multiprocessing.Pool(para)


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 numpy_apply_policy(x, policy):
    return augmentations.apply_policy(policy, x).astype('f')


def stack_augment(augment: list):
    def func(x):
        xl = [augment[i](x) if augment[i] is not None else x for i in range(len(augment))]
        return {k: tf.stack([x[k] for x in xl]) for k in xl[0].keys()}

    return func


class Primitives:
    @staticmethod
    def m():
        return lambda x: augment_mirror(x['image'])

    @staticmethod
    def ms(shift):
        return lambda x: augment_shift(augment_mirror(x['image']), shift)

    @staticmethod
    def s(shift):
        return lambda x: augment_shift(x['image'], shift)


AugmentPair = collections.namedtuple('AugmentPair', 'tf numpy')
PoolEntry = collections.namedtuple('PoolEntry', 'payload batch')


class AugmentPool:
    def __init__(self, get_samples):
        self.get_samples = get_samples

    def __call__(self, *args, **kwargs):
        return self.get_samples()


NOAUGMENT = AugmentPair(tf=lambda x: dict(image=x['image'], label=x['label'], index=x.get('index', -1)),
                        numpy=AugmentPool)


class AugmentPoolAA(AugmentPool):

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

    @staticmethod
    def numpy_apply_policies(arglist):
        x, policies = arglist
        return np.stack([augmentations.apply_policy(policy, y) for y, policy in zip(x, policies)]).astype('f')

    def queue_images(self, batch):
        args = []
        image = batch['image']
        if image.ndim == 4:
            for x in range(image.shape[0]):
                args.append((image[x:x + 1], [random.choice(POLICIES[self.policy_group])]))
        else:
            for x in image[:, 1:]:
                args.append((x, [random.choice(POLICIES[self.policy_group]) for _ in range(x.shape[0])]))
        self.queue.append(PoolEntry(payload=POOL.imap(self.numpy_apply_policies, args), batch=batch))

    def fill_queue(self):
        for _ in range(4):
            self.queue_images(self.get_samples())

    def __call__(self, *args, **kwargs):
        del args, kwargs
        batch = self.get_samples()
        entry = self.queue.pop(0)
        samples = np.stack(list(entry.payload))
        if entry.batch['image'].ndim == 4:
            samples = samples.reshape(entry.batch['image'].shape)
            entry.batch['image'] = samples
        else:
            samples = samples.reshape(entry.batch['image'][:, 1:].shape)
            entry.batch['image'][:, 1:] = samples
        self.queue_images(batch)
        return entry.batch


class AugmentPoolRA(AugmentPoolAA):
    def __init__(self, get_samples, nops, magnitude, cutout):
        init_pool()
        self.get_samples = get_samples
        self.nops = nops
        self.magnitude = magnitude
        self.size = cutout
        self.queue = []
        self.fill_queue()

    @staticmethod
    def numpy_apply_policies(arglist):
        x, policies, cutout = arglist
        return np.stack([augmentations.cutout_numpy(augmentations.apply_policy(policy, y),
                                                    size=int(0.01 * cutout * min(y.shape[:2])))
                         for y, policy in zip(x, policies)]).astype('f')

    def queue_images(self, batch):
        args = []
        image = batch['image']
        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),
                        image=ctaugment.apply(x, cta.policy(probe=False)))
        assert not probe
        return dict(image=np.stack([x[0]] + [ctaugment.apply(y, cta.policy(probe=False)) for y in x[1:]]).astype('f'))

    def queue_images(self):
        batch = self.get_samples()
        args = [(x, batch['cta'], batch['probe']) for x in batch['image']]
        self.queue.append(PoolEntry(payload=POOL.imap(self.numpy_apply_policies, args), batch=batch))

    def fill_queue(self):
        for _ in range(4):
            self.queue_images()

    def __call__(self, *args, **kwargs):
        del args, kwargs
        entry = self.queue.pop(0)
        samples = list(entry.payload)
        entry.batch['image'] = np.stack(x['image'] for x in samples)
        if 'probe' in samples[0]:
            entry.batch['probe'] = np.stack(x['probe'] for x in samples)
            entry.batch['policy'] = [x['policy'] for x in samples]
        self.queue_images()
        return entry.batch


DEFAULT_AUGMENT = EasyDict(
    cifar10=AugmentPair(tf=lambda x: dict(image=Primitives.ms(4)(x), label=x['label'], index=x.get('index', -1)),
                        numpy=AugmentPool),
    cifar100=AugmentPair(tf=lambda x: dict(image=Primitives.ms(4)(x), label=x['label'], index=x.get('index', -1)),
                         numpy=AugmentPool),
    fashion_mnist=AugmentPair(tf=lambda x: dict(image=Primitives.ms(4)(x), label=x['label'], index=x.get('index', -1)),
                              numpy=AugmentPool),
    stl10=AugmentPair(tf=lambda x: dict(image=Primitives.ms(12)(x), label=x['label'], index=x.get('index', -1)),
                      numpy=AugmentPool),
    svhn=AugmentPair(tf=lambda x: dict(image=Primitives.s(4)(x), label=x['label'], index=x.get('index', -1)),
                     numpy=AugmentPool),
    svhn_noextra=AugmentPair(tf=lambda x: dict(image=Primitives.s(4)(x), label=x['label'], index=x.get('index', -1)),
                             numpy=AugmentPool),
)
AUTO_AUGMENT = EasyDict({
    k: AugmentPair(tf=v.tf, numpy=functools.partial(AugmentPoolAA, policy_group=k))
    for k, v in DEFAULT_AUGMENT.items()
})


def get_augmentation(dataset: str, augmentation: str):
    if augmentation == 'x':
        return NOAUGMENT
    elif augmentation == 'm':
        return AugmentPair(tf=lambda x: dict(image=Primitives.m()(x), label=x['label'], index=x.get('index', -1)),
                           numpy=AugmentPool)
    elif augmentation == 'd':
        return DEFAULT_AUGMENT[dataset]
    elif augmentation == 'aa':
        return AUTO_AUGMENT[dataset]
    elif augmentation[0] == 'r':
        nops, mag, cutout = (int(x) for x in augmentation[1:].split('_'))
        return AugmentPair(tf=DEFAULT_AUGMENT[dataset].tf,
                           numpy=functools.partial(AugmentPoolRA, nops=nops, magnitude=mag, cutout=cutout))
    else:
        raise NotImplementedError(augmentation)


def augment_function(dataset: str):
    augmentations = FLAGS.augment.split('.')
    assert len(augmentations) == 2
    return [get_augmentation(dataset, x) for x in augmentations]


def pair_augment_function(dataset: str):
    augmentations = FLAGS.augment.split('.')
    assert len(augmentations) == 3
    unlabeled = [get_augmentation(dataset, x) for x in augmentations[1:]]
    return [get_augmentation(dataset, augmentations[0]),
            AugmentPair(tf=stack_augment([x.tf for x in unlabeled]), numpy=unlabeled[-1].numpy)]


def many_augment_function(dataset: str):
    augmentations = FLAGS.augment.split('.')
    assert len(augmentations) == 3
    unlabeled = [get_augmentation(dataset, x) for x in (augmentations[1:2] + augmentations[2:] * FLAGS.K)]
    return [get_augmentation(dataset, augmentations[0]),
            AugmentPair(tf=stack_augment([x.tf for x in unlabeled]), numpy=unlabeled[-1].numpy)]