Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from keras.objectives import *
- import keras_contrib.backend as KC
- class DSSIM_TD():
- def __init__(self, k1=0.01, k2=0.03, kernel_size=3, max_value=1.0,n_imgs=0):
- """
- A Timedistributed Version Difference of Structural Similarity (DSSIM loss function). Clipped between 0 and 0.5
- Note : You should add a regularization term like a l2 loss in addition to this one.
- Note : In theano, the `kernel_size` must be a factor of the output size. So 3 could
- not be the `kernel_size` for an output of 32.
- # Arguments
- k1: Parameter of the SSIM (default 0.01)
- k2: Parameter of the SSIM (default 0.03)
- kernel_size: Size of the sliding window (default 3)
- max_value: Max value of the output (default 1.0)
- n_imgs: Number of prediected Images in each sample
- """
- self.__name__ = 'DSSIMObjective_TD'
- self.kernel_size = kernel_size
- self.k1 = k1
- self.k2 = k2
- self.max_value = max_value
- self.c1 = (self.k1 * self.max_value) ** 2
- self.c2 = (self.k2 * self.max_value) ** 2
- self.dim_ordering = K.image_data_format()
- self.backend = KC.backend()
- self.n_imgs=n_imgs
- if n_imgs==0:
- raise ValueError('n_imgs can\'t equal zero,use DSSIMObjective instead')
- def __int_shape(self, x):
- return KC.int_shape(x) if self.backend == 'tensorflow' else KC.shape(x)
- def __call__(self, y_true_o, y_pred_o):
- # There are additional parameters for this function
- # Note: some of the 'modes' for edge behavior do not yet have a gradient definition in the Theano tree
- # and cannot be used for learning
- kernel = [self.kernel_size, self.kernel_size]
- ks=[]
- total_t=self.max_value*self.n_imgs
- for in_image in range(0,self.n_imgs):
- y_true=y_true_o[:,in_image]
- y_pred=y_pred_o[:,in_image]
- y_true = KC.reshape(y_true, [-1] + list(self.__int_shape(y_pred)[1:]))
- y_pred = KC.reshape(y_pred, [-1] + list(self.__int_shape(y_pred)[1:]))
- patches_pred = KC.extract_image_patches(y_pred, kernel, kernel, 'valid', self.dim_ordering)
- patches_true = KC.extract_image_patches(y_true, kernel, kernel, 'valid', self.dim_ordering)
- # Reshape to get the var in the cells
- bs, w, h, c1, c2, c3 = self.__int_shape(patches_pred)
- patches_pred = KC.reshape(patches_pred, [-1, w, h, c1 * c2 * c3])
- patches_true = KC.reshape(patches_true, [-1, w, h, c1 * c2 * c3])
- # Get mean
- u_true = KC.mean(patches_true, axis=-1)
- u_pred = KC.mean(patches_pred, axis=-1)
- # Get variance
- var_true = K.var(patches_true, axis=-1)
- var_pred = K.var(patches_pred, axis=-1)
- # Get std dev
- covar_true_pred = K.mean(patches_true * patches_pred, axis=-1) - u_true * u_pred
- ssim = (2 * u_true * u_pred + self.c1) * (2 * covar_true_pred + self.c2)
- denom = (K.square(u_true) + K.square(u_pred) + self.c1) * (var_pred + var_true + self.c2)
- ssim /= denom # no need for clipping, c1 and c2 make the denom non-zero
- ks.append(K.mean((1.0 - ssim) / 2.0))
- all_summed=0
- for i in ks:
- all_summed+=i
- return all_summed/total_t
Add Comment
Please, Sign In to add comment