Search code examples
pythontensorflow

How to realise the 2-D Gaussian filter (like the scipy.ndimage.gaussian_filter) in TensorFlow?


I am just trying to make a Gaussian filter (like the 'scipy.ndimage.gaussian_filter') to process a 4-D tensor in TensorFlow, the 4-D tensor has a shape of: [16,96,96,3] (16 is the batch size, 96 is the image block size, and 3 is the number of channels). How could I achieve this?


Solution

  • You just need to create a Gaussian 2D kernel and use a 2D convolution:

    import tensorflow as tf
    
    # Make Gaussian kernel following SciPy logic
    def make_gaussian_2d_kernel(sigma, truncate=4.0, dtype=tf.float32):
        radius = tf.to_int32(sigma * truncate)
        x = tf.cast(tf.range(-radius, radius + 1), dtype=dtype)
        k = tf.exp(-0.5 * tf.square(x / sigma))
        k = k / tf.reduce_sum(k)
        return tf.expand_dims(k, 1) * k
    
    # Input data
    image = tf.placeholder(tf.float32, [16, 96, 96, 3])
    # Convolution kernel
    kernel = make_gaussian_2d_kernel(5)
    # Apply kernel to each channel (see https://stackoverflow.com/q/55687616/1782792)
    kernel = tf.tile(kernel[:, :, tf.newaxis, tf.newaxis], [1, 1, 3, 1])
    image_filtered = tf.nn.separable_conv2d(
        image, kernel, tf.eye(3, batch_shape=[1, 1]),
        strides=[1, 1, 1, 1], padding='SAME')