tf_bilateral_rgb.py

tf_bilateral_rgb.py

    #!/usr/bin/python
# torch_bilateral: bi/trilateral filtering in tf
# Copyright: (c) 2017-2018 Prophesee
from __future__ import print_function
import tensorflow as tf
import numpy as np
import pdb
import time


def gkern2d(l=21, sig=3):
    """Returns a 2D Gaussian kernel array."""
    ax = np.arange(-l // 2 + 1., l // 2 + 1.)
    xx, yy = np.meshgrid(ax, ax)
    kernel = np.exp(-(xx ** 2 + yy ** 2) / (2. * sig ** 2))
    return kernel



def shift(image,kernel_size,w,h):
    if kernel_size == 3:
        pad = 1
    elif kernel_size == 5:
        pad = 2
    elif kernel_size == 7:
        pad = 3

    paddings = tf.constant([[0,0], [0,0], [0,0], [pad, pad], [pad, pad]])
    x_pad = tf.pad(image, paddings, "CONSTANT") # "REFLECT"

    #gather channels
    cat_layers = []
    for y in range(0, kernel_size):
        y2 = y + h
        for x in range(0, kernel_size):
            x2 = x + w
            crop = x_pad[:,:,:,y:y2,x:x2]
            cat_layers.append(crop)

    y = tf.concat(cat_layers,2)

    return y

def bilateral(image,kernel_size,channels,height,width,sigma_space=3,sigma_color=0.1):
    gkern = gkern2d(kernel_size, sigma_space).reshape(1,1,kernel_size**2,1,1).repeat(channels,1)
    gw =  tf.convert_to_tensor(gkern, dtype=tf.float32)

    input = tf.expand_dims(image, 2)
    ishift = shift(input,kernel_size,w,h)
    iex = tf.tile(input,[1,1,kernel_size**2,1,1])
    dist = tf.pow(ishift - iex,2)
    gauss = tf.exp(-dist / sigma_color) * gw
    wsum = tf.reduce_sum(gauss, axis=2)
    ifilt = tf.reduce_sum(gauss * ishift, axis=2) / wsum
    return ifilt



if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import cv2
    import os

    c, h, w = 3, 480, 640
    k = 5

    #os.environ['CUDA_VISIBLE_DEVICES'] = ''

    #Define the Graph
    with tf.device('/device:GPU:0'): #
        input = tf.placeholder(tf.float32, shape=[None, c, h, w], name='image_input')
        output = bilateral(input,k,c,h,w)


    cuda = False

    im = cv2.imread('/home/eperot/Pictures/Lena.png', cv2.IMREAD_GRAYSCALE if c == 1 else cv2.IMREAD_COLOR)
    im = cv2.resize(im,(w,h),interpolation=cv2.INTER_CUBIC)


    im_in = np.moveaxis(im, 2, 0)
    im_in = im_in.reshape(1, c, h, w).astype(np.float32) / 255.0


    with tf.Session() as sess:

        start = time.time()
        out = sess.run([output], feed_dict={input: im_in})[0]
        print(time.time()-start, ' s')

        img_out = np.moveaxis(out[0], 0, 2) * 255

        print(img_out.min(),img_out.max())

        cv2.imshow('img_in', im.astype(np.uint8))
        cv2.imshow('img',img_out.astype(np.uint8))
        cv2.waitKey()