|
- import numpy as np
- from PIL import Image
- import tensorflow as tf
- import scipy.stats as st
- from skimage import io,data,color
- from functools import reduce
-
- def gauss_kernel(kernlen=21, nsig=3, channels=1):
- interval = (2*nsig+1.)/(kernlen)
- x = np.linspace(-nsig-interval/2., nsig+interval/2., kernlen+1)
- kern1d = np.diff(st.norm.cdf(x))
- kernel_raw = np.sqrt(np.outer(kern1d, kern1d))
- kernel = kernel_raw/kernel_raw.sum()
- out_filter = np.array(kernel, dtype = np.float32)
- out_filter = out_filter.reshape((kernlen, kernlen, 1, 1))
- out_filter = np.repeat(out_filter, channels, axis = 2)
- return out_filter
-
- def tensor_size(tensor):
- from operator import mul
- return reduce(mul, (d.value for d in tensor.get_shape()[1:]), 1)
-
- def blur(x):
- kernel_var = gauss_kernel(21, 3, 3)
- return tf.nn.depthwise_conv2d(x, kernel_var, [1, 1, 1, 1], padding='SAME')
-
- def tensor_size(tensor):
- from operator import mul
- return reduce(mul, (d.value for d in tensor.get_shape()[1:]), 1)
-
- def data_augmentation(image, mode):
- if mode == 0:
- # original
- return image
- elif mode == 1:
- # flip up and down
- return np.flipud(image)
- elif mode == 2:
- # rotate counterwise 90 degree
- return np.rot90(image)
- elif mode == 3:
- # rotate 90 degree and flip up and down
- image = np.rot90(image)
- return np.flipud(image)
- elif mode == 4:
- # rotate 180 degree
- return np.rot90(image, k=2)
- elif mode == 5:
- # rotate 180 degree and flip
- image = np.rot90(image, k=2)
- return np.flipud(image)
- elif mode == 6:
- # rotate 270 degree
- return np.rot90(image, k=3)
- elif mode == 7:
- # rotate 270 degree and flip
- image = np.rot90(image, k=3)
- return np.flipud(image)
-
- def load_images(file):
- im = Image.open(file)
- img = np.array(im, dtype="float32") / 255.0
- img_max = np.max(img)
- img_min = np.min(img)
- img_norm = np.float32((img - img_min) / np.maximum((img_max - img_min), 0.001))
- return img_norm
-
- def gradient(input_tensor, direction):
- smooth_kernel_x = tf.reshape(tf.constant([[0, 0], [-1, 1]], tf.float32), [2, 2, 1, 1])
- smooth_kernel_y = tf.transpose(smooth_kernel_x, [1, 0, 2, 3])
- if direction == "x":
- kernel = smooth_kernel_x
- elif direction == "y":
- kernel = smooth_kernel_y
- gradient_orig = tf.abs(tf.nn.conv2d(input_tensor, kernel, strides=[1, 1, 1, 1], padding='SAME'))
- grad_min = tf.reduce_min(gradient_orig)
- grad_max = tf.reduce_max(gradient_orig)
- grad_norm = tf.div((gradient_orig - grad_min), (grad_max - grad_min + 0.0001))
- return grad_norm
-
- def _tf_fspecial_gauss(size, sigma):
- """Function to mimic the 'fspecial' gaussian MATLAB function
- """
- x_data, y_data = np.mgrid[-size//2 + 1:size//2 + 1, -size//2 + 1:size//2 + 1]
-
- x_data = np.expand_dims(x_data, axis=-1)
- x_data = np.expand_dims(x_data, axis=-1)
-
- y_data = np.expand_dims(y_data, axis=-1)
- y_data = np.expand_dims(y_data, axis=-1)
-
- x = tf.constant(x_data, dtype=tf.float32)
- y = tf.constant(y_data, dtype=tf.float32)
-
- g = tf.exp(-((x**2 + y**2)/(2.0*sigma**2)))
- return g / tf.reduce_sum(g)
-
- def tf_ssim(img1, img2, cs_map=False, mean_metric=True, size=11, sigma=1.5):
- window = _tf_fspecial_gauss(size, sigma) # window shape [size, size]
- K1 = 0.01
- K2 = 0.03
- L = 1 # depth of image (255 in case the image has a differnt scale)
- C1 = (K1*L)**2
- C2 = (K2*L)**2
- mu1 = tf.nn.conv2d(img1, window, strides=[1,1,1,1], padding='VALID')
- mu2 = tf.nn.conv2d(img2, window, strides=[1,1,1,1],padding='VALID')
- mu1_sq = mu1*mu1
- mu2_sq = mu2*mu2
- mu1_mu2 = mu1*mu2
- sigma1_sq = tf.nn.conv2d(img1*img1, window, strides=[1,1,1,1],padding='VALID') - mu1_sq
- sigma2_sq = tf.nn.conv2d(img2*img2, window, strides=[1,1,1,1],padding='VALID') - mu2_sq
- sigma12 = tf.nn.conv2d(img1*img2, window, strides=[1,1,1,1],padding='VALID') - mu1_mu2
- if cs_map:
- value = (((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*
- (sigma1_sq + sigma2_sq + C2)),
- (2.0*sigma12 + C2)/(sigma1_sq + sigma2_sq + C2))
- else:
- value = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*
- (sigma1_sq + sigma2_sq + C2))
-
- if mean_metric:
- value = tf.reduce_mean(value)
- return value
-
- def save_images(filepath, result_1, result_2 = None, result_3 = None):
- result_1 = np.squeeze(result_1)
- result_2 = np.squeeze(result_2)
- result_3 = np.squeeze(result_3)
-
- if not result_2.any():
- cat_image = result_1
- else:
- cat_image = np.concatenate([result_1, result_2], axis = 1)
- if not result_3.any():
- cat_image = cat_image
- else:
- cat_image = np.concatenate([cat_image, result_3], axis = 1)
-
- im = Image.fromarray(np.clip(cat_image * 255.0, 0, 255.0).astype('uint8'))
- im.save(filepath, 'png')
|