LOSS.py

import tensorflow as tf
import numpy as np


def SSIM_LOSS(img1, img2, 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 different 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

	value = (2.0 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2)
	value = tf.reduce_mean(value)
	return value


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 L1_LOSS(batchimg):
	L1_norm = tf.reduce_sum(tf.abs(batchimg), axis = [1, 2])
	# tf.norm(batchimg, axis = [1, 2], ord = 1) / int(batchimg.shape[1])
	E = tf.reduce_mean(L1_norm)
	return E


def Fro_LOSS(batchimg):
	fro_norm = tf.square(tf.norm(batchimg, axis = [1, 2], ord = 'fro'))
		# / (int(batchimg.shape[1]) * int(batchimg.shape[2]))
	E = tf.reduce_mean(fro_norm)
	return E