Gaussian Blur in Tensorflow.JS

gaussianBlur.ts

import * as tf from '@tensorflow/tfjs';
import { Tensor3D, Tensor2D, Tensor1D, Tensor4D, Tensor } from '@tensorflow/tfjs';


function get1dGaussianKernel(sigma, size): Tensor1D {
    // Generate a 1d gaussian distribution across a range
    var x = tf.range(Math.floor(-size / 2) + 1, Math.floor(size / 2) + 1)
    x = tf.pow(x, 2)
    x = tf.exp(x.div(-2.0 * (sigma * sigma))) as Tensor1D
    x = x.div(tf.sum(x))
    return x
}

function get2dGaussianKernel(size: number, sigma?: number): Tensor2D {
    // This default is to mimic opencv2. 
    sigma = sigma || (0.3 * ((size - 1) * 0.5 - 1) + 0.8)

    var kerne1d = get1dGaussianKernel(sigma, size)
    return tf.outerProduct(kerne1d, kerne1d)
}

export function getGaussianKernel(size = 5, sigma?: number): Tensor4D {
    return tf.tidy(() => {
        var kerne2d = get2dGaussianKernel(size, sigma)
        var kerne3d = tf.stack([kerne2d, kerne2d, kerne2d])
        return tf.reshape(kerne3d, [size, size, 3, 1])
    })
}

export function blur(image: Tensor3D, kernel: Tensor4D): Tensor3D {
    return tf.tidy(() => {
        return tf.depthwiseConv2d(image, kernel, 1, "valid")
    })
}

tensorflow-js-blur-example.ts

import { getGaussianKernel, blur } from "./gaussianBlur"

// Create a 7x7 Gaussian Kernel
const kernel = getGaussianKernel(7)

// apply kernel on a image tensor shaped [height, width, 3]
var blurredImg = blur(img, kernel)