subpixelupscaling.py

# -*- coding: utf-8 -*-
"""
Created on Sat Jan 18 21:24:15 2020

@author: sabari
"""

# -*- coding: utf-8 -*-
from __future__ import absolute_import

from keras.layers import Layer
import tensorflow as tf
from keras_contrib import backend as KC
from keras_contrib.utils.conv_utils import normalize_data_format
from keras.layers import Lambda

class SubPixelUpscaling(Layer):
    """ Sub-pixel convolutional upscaling layer.
    This layer requires a Convolution2D prior to it,
    having output filters computed according to
    the formula :
        filters = k * (scale_factor * scale_factor)
        where k = a user defined number of filters (generally larger than 32)
              scale_factor = the upscaling factor (generally 2)
    This layer performs the depth to space operation on
    the convolution filters, and returns a
    tensor with the size as defined below.
    # Example :
    ```python
        # A standard subpixel upscaling block
        x = Convolution2D(256, 3, 3, padding='same', activation='relu')(...)
        u = SubPixelUpscaling(scale_factor=2)(x)
        # Optional
        x = Convolution2D(256, 3, 3, padding='same', activation='relu')(u)
    ```
    In practice, it is useful to have a second convolution layer after the
    SubPixelUpscaling layer to speed up the learning process.
    However, if you are stacking multiple
    SubPixelUpscaling blocks, it may increase
    the number of parameters greatly, so the
    Convolution layer after SubPixelUpscaling
    layer can be removed.
    # Arguments
        scale_factor: Upscaling factor.
        data_format: Can be None, 'channels_first' or 'channels_last'.
    # Input shape
        4D tensor with shape:
        `(samples, k * (scale_factor * scale_factor) channels, rows, cols)`
        if data_format='channels_first'
        or 4D tensor with shape:
        `(samples, rows, cols, k * (scale_factor * scale_factor) channels)`
        if data_format='channels_last'.
    # Output shape
        4D tensor with shape:
        `(samples, k channels, rows * scale_factor, cols * scale_factor))`
        if data_format='channels_first'
        or 4D tensor with shape:
        `(samples, rows * scale_factor, cols * scale_factor, k channels)`
        if data_format='channels_last'.
    # References
        - [Real-Time Single Image and Video Super-Resolution Using an
           Efficient Sub-Pixel Convolutional Neural Network](
           https://arxiv.org/abs/1609.05158)
    """

    def __init__(self, scale_factor=2, data_format=None, **kwargs):
        super(SubPixelUpscaling, self).__init__(**kwargs)

        self.scale_factor = scale_factor
        self.data_format = normalize_data_format(data_format)

    def build(self, input_shape):
        pass

    def call(self, x, mask=None):
        y = KC.depth_to_space(x, self.scale_factor, self.data_format)
        return y

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_first':
            b, k, r, c = input_shape
            new_k = k // (self.scale_factor ** 2)
            new_r = r * self.scale_factor
            new_c = c * self.scale_factor
            return b, new_k, new_r, new_c
        else:
            b, r, c, k = input_shape
            new_r = r * self.scale_factor
            new_c = c * self.scale_factor
            new_k = k // (self.scale_factor ** 2)
            return b, new_r, new_c, new_k

    def get_config(self):
        config = {'scale_factor': self.scale_factor,
                  'data_format': self.data_format}
        base_config = super(SubPixelUpscaling, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
    
    
def SubpixelConv2D(input_shape, scale=4):
    """
    Keras layer to do subpixel convolution.
    NOTE: Tensorflow backend only. Uses tf.depth_to_space
    Ref:
        [1] Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network
            Shi et Al.
            https://arxiv.org/abs/1609.05158
    :param input_shape: tensor shape, (batch, height, width, channel)
    :param scale: upsampling scale. Default=4
    :return:
    """
    # upsample using depth_to_space
    def subpixel_shape(input_shape):
        dims = [input_shape[0],
                input_shape[1] * scale,
                input_shape[2] * scale,
                int(input_shape[3] / (scale ** 2))]
        output_shape = tuple(dims)
        return output_shape

    def subpixel(x):
        return tf.depth_to_space(x, scale)


    return Lambda(subpixel, output_shape=subpixel_shape)