layers.py

'''
Custom Keras layers used on the pastiche model.
'''

import tensorflow as tf
import keras
from keras import initializations
from keras.layers import ZeroPadding2D, Layer, InputSpec

# Extending the ZeroPadding2D layer to do reflection padding instead.
class ReflectionPadding2D(ZeroPadding2D):
    def call(self, x, mask=None):
        pattern = [[0, 0],
                   [self.top_pad, self.bottom_pad],
                   [self.left_pad, self.right_pad],
                   [0, 0]]
        return tf.pad(x, pattern, mode='REFLECT')


class InstanceNormalization(Layer):
    def __init__(self, epsilon=1e-5, weights=None,
                 beta_init='zero', gamma_init='one', **kwargs):
        self.beta_init = initializations.get(beta_init)
        self.gamma_init = initializations.get(gamma_init)
        self.epsilon = epsilon
        super(InstanceNormalization, self).__init__(**kwargs)

    def build(self, input_shape):
        # This currently only works for 4D inputs: assuming (B, H, W, C)
        self.input_spec = [InputSpec(shape=input_shape)]
        shape = (1, 1, 1, input_shape[-1])

        self.gamma = self.gamma_init(shape, name='{}_gamma'.format(self.name))
        self.beta = self.beta_init(shape, name='{}_beta'.format(self.name))
        self.trainable_weights = [self.gamma, self.beta]

        self.built = True

    def call(self, x, mask=None):
        # Do not regularize batch axis
        reduction_axes = [1, 2]

        mean, var = tf.nn.moments(x, reduction_axes,
                                  shift=None, name=None, keep_dims=True)
        x_normed = tf.nn.batch_normalization(x, mean, var, self.beta, self.gamma, self.epsilon)
        return x_normed

    def get_config(self):
        config = {"epsilon": self.epsilon}
        base_config = super(InstanceNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConditionalInstanceNormalization(InstanceNormalization):
    def __init__(self, targets, nb_classes, **kwargs):
        self.targets = targets
        self.nb_classes = nb_classes
        super(ConditionalInstanceNormalization, self).__init__(**kwargs)

    def build(self, input_shape):
        # This currently only works for 4D inputs: assuming (B, H, W, C)
        self.input_spec = [InputSpec(shape=input_shape)]
        shape = (self.nb_classes, 1, 1, input_shape[-1])

        self.gamma = self.gamma_init(shape, name='{}_gamma'.format(self.name))
        self.beta = self.beta_init(shape, name='{}_beta'.format(self.name))
        self.trainable_weights = [self.gamma, self.beta]

        self.built = True

    def call(self, x, mask=None):
        # Do not regularize batch axis
        reduction_axes = [1, 2]

        mean, var = tf.nn.moments(x, reduction_axes,
                                  shift=None, name=None, keep_dims=True)

        # Get the appropriate lines of gamma and beta
        beta = tf.gather(self.beta, self.targets)
        gamma = tf.gather(self.gamma, self.targets)
        x_normed = tf.nn.batch_normalization(x, mean, var, beta, gamma, self.epsilon)

        return x_normed