custom_layer.py

import tensorflow as tf

def tf_int_round(num):
	return tf.cast(tf.round(num), dtype=tf.int32)

class resize_layer(layers.Layer):
    # Initialize variables
    def __init__(self, scale, **kwargs):
        self.scale = scale
        super(resize_layer, self).__init__(**kwargs)
        
    def build(self, input_shape):
        super(resize_layer,self).build(input_shape)
        
    # Defining how we will call our function    
    def call(self, x, method="bicubic"):
    	height = tf_int_round(tf.cast(tf.shape(x)[1],dtype=tf.float32) * self.scale)
    	width = tf_int_round(tf.cast(tf.shape(x)[2],dtype=tf.float32) * self.scale)

    	if method == "bilinear":
    		return tf.image.resize_bilinear(x, size=(height, width))
        elif method == "bicubic":
        	return tf.image.resize_bicubic(x, size=(height, width))
        elif method == "nearest":
        	return tf.image.resize_nearest_neighbor(x, size=(height, width))
        
    # Defining the computation of the output shape
    def get_output_shape_for(self, input_shape):
    	height = tf_int_round(tf.cast(tf.shape(x)[1],dtype=tf.float32) * self.scale)
    	width = tf_int_round(tf.cast(tf.shape(x)[2],dtype=tf.float32) * self.scale)
        return (self.input_shape[0], height, width, input_shape[3])

# Using our new custom layer with the Functional API 
image_2 = resize_layer(scale=2)(image, method="bilinear")