ops.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


import tensorflow as tf
import tensorflow.contrib.slim as slim
import pdb


def preprocess(image):
    with tf.name_scope("preprocess"):
        # [0, 1] => [-1, 1]
        return image * 2 - 1


def deprocess(image):
    with tf.name_scope("deprocess"):
        # [-1, 1] => [0, 1]
        return (image + 1) / 2


def preprocessLR(image):
    with tf.name_scope("preprocessLR"):
        return tf.identity(image)


def deprocessLR(image):
    with tf.name_scope("deprocessLR"):
        return tf.identity(image)


# Define the convolution building block
def conv2(batch_input, kernel=3, output_channel=64, stride=1, use_bias=True, scope='conv',norm=None):
    # kernel: An integer specifying the width and height of the 2D convolution window
    with tf.variable_scope(scope):     
        if norm == None:
            if use_bias:
                return slim.conv2d(batch_input, output_channel, [kernel, kernel], stride, 'SAME', data_format='NHWC',
                                activation_fn=None, weights_initializer=tf.contrib.layers.xavier_initializer())
            else:
                return slim.conv2d(batch_input, output_channel, [kernel, kernel], stride, 'SAME', data_format='NHWC',
                                activation_fn=None, weights_initializer=tf.contrib.layers.xavier_initializer(),
                                biases_initializer=None)
        elif norm == 'weight_norm':
            w = tf.get_variable("kernel", shape=[kernel, kernel, batch_input.get_shape()[-1], output_channel],                                                                    initializer=tf.contrib.layers.xavier_initializer())
            x = tf.nn.conv2d(input=batch_input, filter=weight_norm(w, axis=[0,1,2]),strides=[1, stride, stride, 1], padding='SAME')
            if use_bias:
                bias = tf.get_variable("bias", [output_channel], initializer=tf.constant_initializer(0.0))
                x = tf.nn.bias_add(x, bias)
                
            return x
                
        elif norm == 'spectral_norm': 
            w = tf.get_variable("kernel", shape=[kernel, kernel, batch_input.get_shape()[-1], output_channel],                                                                    initializer=tf.contrib.layers.xavier_initializer())
            x = tf.nn.conv2d(input=batch_input, filter=spectral_norm(w),strides=[1, stride, stride, 1], padding='SAME')
            if use_bias:
                bias = tf.get_variable("bias", [output_channel], initializer=tf.constant_initializer(0.0))
                x = tf.nn.bias_add(x, bias)
                
            return x


def conv2_NCHW(batch_input, kernel=3, output_channel=64, stride=1, use_bias=True, scope='conv_NCHW'):
    # Use NCWH to speed up the inference
    # kernel: list of 2 integer specifying the width and height of the 2D convolution window
    with tf.variable_scope(scope):
        if use_bias:
            return slim.conv2d(batch_input, output_channel, [kernel, kernel], stride, 'SAME', data_format='NCWH',
                               activation_fn=None, weights_initializer=tf.contrib.layers.xavier_initializer())
        else:
            return slim.conv2d(batch_input, output_channel, [kernel, kernel], stride, 'SAME', data_format='NCWH',
                               activation_fn=None, weights_initializer=tf.contrib.layers.xavier_initializer(),
                               biases_initializer=None)


# Define our tensorflow version PRelu
def prelu_tf(inputs, name='Prelu'):
    with tf.variable_scope(name):
        alphas = tf.get_variable('alpha', inputs.get_shape()[-1], initializer=tf.zeros_initializer(), dtype=tf.float32)
    pos = tf.nn.relu(inputs)
    neg = alphas * (inputs - abs(inputs)) * 0.5

    return pos + neg


# Define our Lrelu
def lrelu(inputs, alpha):
    return tf.keras.layers.LeakyReLU(alpha=alpha).call(inputs)


def batchnorm(inputs, is_training):
    return slim.batch_norm(inputs, decay=0.9, epsilon=0.001, updates_collections=tf.GraphKeys.UPDATE_OPS,
                        scale=False, fused=True, is_training=is_training)


# Our dense layer
def denselayer(inputs, output_size,norm=None):
    x = flatten(inputs)
    shape = x.get_shape().as_list()
    channels = shape[-1]
    
    if norm == None: 
        output = tf.layers.dense(inputs, output_size, activation=None, kernel_initializer=tf.contrib.layers.xavier_initializer())
    elif norm == 'weight_norm':
        w = tf.get_variable("kernel", [channels, units], tf.float32,initializer=weight_init, regularizer=weight_regularizer)
        bias = tf.get_variable("bias", [units],initializer=tf.constant_initializer(0.0))
        output = tf.matmul(x, weight_norm(w)) + bias
        
    elif norm == 'spectral_norm':
        w = tf.get_variable("kernel", [channels, units], tf.float32,initializer=weight_init, regularizer=weight_regularizer)
        bias = tf.get_variable("bias", [units],initializer=tf.constant_initializer(0.0))
        output = tf.matmul(x, spectral_norm(w)) + bias
        
    return output

def flatten(x) :
    return tf.layers.flatten(x)

def hw_flatten(x) :
    return tf.reshape(x, shape=[x.shape[0], -1, x.shape[-1]])

# The implementation of PixelShuffler
def pixelShuffler(inputs, scale=2):
    size = tf.shape(inputs)
    batch_size = size[0]
    h = size[1]
    w = size[2]
    c = inputs.get_shape().as_list()[-1]

    # Get the target channel size
    channel_target = c // (scale * scale)
    channel_factor = c // channel_target

    shape_1 = [batch_size, h, w, channel_factor // scale, channel_factor // scale]
    shape_2 = [batch_size, h * scale, w * scale, 1]

    # Reshape and transpose for periodic shuffling for each channel
    input_split = tf.split(inputs, channel_target, axis=3)
    output = tf.concat([phaseShift(x, scale, shape_1, shape_2) for x in input_split], axis=3)

    return output


def phaseShift(inputs, scale, shape_1, shape_2):
    # Tackle the condition when the batch is None
    X = tf.reshape(inputs, shape_1)
    X = tf.transpose(X, [0, 1, 3, 2, 4])

    return tf.reshape(X, shape_2)


# The random flip operation used for loading examples
def random_flip(input, decision):
    f1 = tf.identity(input)
    f2 = tf.image.flip_left_right(input)
    output = tf.cond(tf.less(decision, 0.5), lambda: f2, lambda: f1)

    return output


# The operation used to print out the configuration
def print_configuration_op(FLAGS):
    print('[Configurations]:')
    a = FLAGS.mode
    #pdb.set_trace()
    for name, value in FLAGS.__flags.items():
        if type(value) == float:
            print('\t%s: %f'%(name, value))
        elif type(value) == int:
            print('\t%s: %d'%(name, value))
        elif type(value) == str:
            print('\t%s: %s'%(name, value))
        elif type(value) == bool:
            print('\t%s: %s'%(name, value))
        else:
            print('\t%s: %s' % (name, value))

    print('End of configuration')


def compute_psnr(ref, target):
    psnr = tf.image.psnr(ref, target, max_val=255)
    return psnr[0]

def compute_psnr_test(ref, target):
    target_shape = tf.shape(target)
    ref_ = tf.image.resize_image_with_crop_or_pad(ref,target_height=target_shape[1],target_width=target_shape[2])    
    psnr = tf.image.psnr(ref_, target, max_val=255)
    return psnr[0]

def compute_ssim(ref,target):
    ssim = tf.image.ssim(ref,target,max_val=255)
    return ssim[0]

def compute_ssim_test(ref,target):
    target_shape = tf.shape(target)
    ref_ = tf.image.resize_image_with_crop_or_pad(ref,target_height = target_shape[1],target_width = target_shape[2])
    ssim = tf.image.ssim(ref_,target,max_val=255)
    return ssim[0]

def evaluation_metric(psnr_list,ssim_list):
    count = len(psnr_list)
    min_psnr = min(psnr_list)
    max_psnr = max(psnr_list)
    mean_psnr = sum(psnr_list)/len(psnr_list)
    min_ssim = min(ssim_list)
    max_ssim = max(ssim_list)
    mean_ssim = sum(ssim_list)/len(ssim_list)

    stats_psnr = {'count': count,'min_psnr': min_psnr, 'max_psnr': max_psnr, 'mean_psnr': mean_psnr }
    stats_ssim = {'count': count,'min_ssim': min_ssim, 'max_ssim': max_ssim, 'mean_ssim' : mean_ssim }
    return stats_psnr, stats_ssim

#Data independent Weight Normalization 
def weight_norm(v, axis=None, name=None, return_all=True,
                       reuse=None,
                       init=tf.random_normal_initializer(stddev=0.05)):
    #v = tf.get_variable(name=name, shape=shape, dtype=tf.float32, initializer=init)
    g = tf.get_variable(name=name+'g', shape=v.shape[-1], initializer=tf.constant_initializer(1),trainable=False)
    
    inv_norm = tf.rsqrt(tf.reduce_sum(tf.square(v), reduction_indices=axis))
    w = v * g * inv_norm
    return w


#Data independent Spectral Normalization  

def spectral_norm(w, iteration=1):
    w_shape = w.shape.as_list()
    w = tf.reshape(w, [-1, w_shape[-1]])

    u = tf.get_variable("u", [1, w_shape[-1]], initializer=tf.truncated_normal_initializer(), trainable=False)

    u_hat = u
    v_hat = None
    for i in range(iteration):
        """
        power iteration
        Usually iteration = 1 will be enough
        """
        v_ = tf.matmul(u_hat, tf.transpose(w))
        v_hat = l2_norm(v_)

        u_ = tf.matmul(v_hat, w)
        u_hat = l2_norm(u_)

    sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))
    w_norm = w / sigma

    with tf.control_dependencies([u.assign(u_hat)]):
        w_norm = tf.reshape(w_norm, w_shape)

    return w_norm

def l2_norm(v, eps=1e-12):
    return v / (tf.reduce_sum(v ** 2) ** 0.5 + eps)

# VGG19 component
def vgg_arg_scope(weight_decay=0.0005):
  """Defines the VGG arg scope.
  Args:
    weight_decay: The l2 regularization coefficient.
  Returns:
    An arg_scope.
  """
  with slim.arg_scope([slim.conv2d, slim.fully_connected],
                      activation_fn=tf.nn.relu,
                      weights_regularizer=slim.l2_regularizer(weight_decay),
                      biases_initializer=tf.zeros_initializer()):
    with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
      return arg_sc


# VGG19 net
def vgg_19(inputs,
           num_classes=1000,
           is_training=False,
           dropout_keep_prob=0.5,
           spatial_squeeze=True,
           scope='vgg_19',
           reuse = False,
           fc_conv_padding='VALID'):
  """Oxford Net VGG 19-Layers version E Example.
  Note: All the fully_connected layers have been transformed to conv2d layers.
        To use in classification mode, resize input to 224x224.
  Args:
    inputs: a tensor of size [batch_size, height, width, channels].
    num_classes: number of predicted classes.
    is_training: whether or not the model is being trained.
    dropout_keep_prob: the probability that activations are kept in the dropout
      layers during training.
    spatial_squeeze: whether or not should squeeze the spatial dimensions of the
      outputs. Useful to remove unnecessary dimensions for classification.
    scope: Optional scope for the variables.
    fc_conv_padding: the type of padding to use for the fully connected layer
      that is implemented as a convolutional layer. Use 'SAME' padding if you
      are applying the network in a fully convolutional manner and want to
      get a prediction map downsampled by a factor of 32 as an output. Otherwise,
      the output prediction map will be (input / 32) - 6 in case of 'VALID' padding.
  Returns:
    the last op containing the log predictions and end_points dict.
  """
  with tf.variable_scope(scope, 'vgg_19', [inputs], reuse=reuse) as sc:
    end_points_collection = sc.name + '_end_points'
    # Collect outputs for conv2d, fully_connected and max_pool2d.
    with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
                        outputs_collections=end_points_collection):
      net = slim.repeat(inputs, 2, slim.conv2d, 64, 3, scope='conv1', reuse=reuse)
      net = slim.max_pool2d(net, [2, 2], scope='pool1')
      net = slim.repeat(net, 2, slim.conv2d, 128, 3, scope='conv2',reuse=reuse)
      net = slim.max_pool2d(net, [2, 2], scope='pool2')
      net = slim.repeat(net, 4, slim.conv2d, 256, 3, scope='conv3', reuse=reuse)
      net = slim.max_pool2d(net, [2, 2], scope='pool3')
      net = slim.repeat(net, 4, slim.conv2d, 512, 3, scope='conv4',reuse=reuse)
      net = slim.max_pool2d(net, [2, 2], scope='pool4')
      net = slim.repeat(net, 4, slim.conv2d, 512, 3, scope='conv5',reuse=reuse)
      net = slim.max_pool2d(net, [2, 2], scope='pool5')
      # Use conv2d instead of fully_connected layers.
      # Convert end_points_collection into a end_point dict.
      end_points = slim.utils.convert_collection_to_dict(end_points_collection)

      return net, end_points
vgg_19.default_image_size = 224