add dssim objective (#26)

* add dssim

* PEP8-ify

* use backend to get the right shape

* Remove the need of the batch size.

* Fix for theano backend on reshape

* PEP8

* Merge master

* changes according to review

keras_contrib/backend/tensorflow_backend.py

@@ -2,6 +2,7 @@ import tensorflow as tf
 from tensorflow.python.training import moving_averages
 from tensorflow.python.ops import tensor_array_ops
 from tensorflow.python.ops import control_flow_ops
+
 try:
     from tensorflow.python.ops import ctc_ops as ctc
 except ImportError:
@@ -28,7 +29,7 @@ def _preprocess_deconv_output_shape(x, shape, dim_ordering):
         shape = (shape[0],) + tuple(shape[2:]) + (shape[1],)
 
     if shape[0] is None:
-        shape = (tf.shape(x)[0], ) + tuple(shape[1:])
+        shape = (tf.shape(x)[0],) + tuple(shape[1:])
         shape = tf.stack(list(shape))
     return shape
 
@@ -100,9 +101,8 @@ def extract_image_patches(x, ksizes, ssizes, border_mode="same",
                                        padding)
     # Reshaping to fit Theano
     bs, w, h, ch = KTF.int_shape(patches)
-    patches = tf.reshape(patches, [bs, w, h, -1, ch_i])
-    patches = tf.reshape(tf.transpose(patches, [0, 1, 2, 4, 3]),
-                         [bs, w, h, ch_i, ksizes[0], ksizes[1]])
+    patches = tf.reshape(tf.transpose(tf.reshape(patches, [-1, w, h, tf.floordiv(ch, ch_i), ch_i]), [0, 1, 2, 4, 3]),
+                         [-1, w, h, ch_i, ksizes[0], ksizes[1]])
     if dim_ordering == "tf":
         patches = KTF.permute_dimensions(patches, [0, 1, 2, 4, 5, 3])
     return patches

keras_contrib/objectives.py

@@ -1,10 +1,64 @@
 from __future__ import absolute_import
 
-from . import backend as K
-from keras.utils.generic_utils import get_from_module
-
 from keras.objectives import *
 
+import keras_contrib.backend as KC
+
 
 def get(identifier):
     return get_from_module(identifier, globals(), 'objective')
+
+
+class DSSIMObjective():
+    def __init__(self, k1=0.01, k2=0.03, kernel_size=3, max_value=1.0):
+        """
+        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.
+        :param batch_size: Batch size used in the model
+        :param k1: Parameter of the SSIM (default 0.01)
+        :param k2: Parameter of the SSIM (default 0.03)
+        :param kernel_size: Size of the sliding window (default 3)
+        :param max_value: Max value of the output (default 1.0)
+        """
+        self.__name__ = "DSSIMObjective"
+        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_dim_ordering()
+        self.backend = KC.backend()
+
+    def __int_shape(self, x):
+        return KC.int_shape(x) if self.backend == "tensorflow" else KC.shape(x)
+
+    def __call__(self, y_true, y_pred):
+        # 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]
+        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
+        std_true = K.sqrt(var_true + KC.epsilon())
+        std_pred = K.sqrt(var_pred + KC.epsilon())
+
+        ssim = (2 * u_true * u_pred + self.c1) * (2 * std_pred * std_true + 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
+        return K.mean((1.0 - ssim) / 2.0)

tests/keras_contrib/test_objectives.py

@@ -1,11 +1,11 @@
-import pytest
 import numpy as np
-
+import pytest
 from keras import backend as K
+from numpy.testing import assert_allclose
+
 from keras_contrib import backend as KC
 from keras_contrib import objectives
 
-
 allobj = []
 
 
@@ -25,5 +25,32 @@ def test_objective_shapes_2d():
         assert K.eval(objective_output).shape == (6,)
 
 
+def test_dssim_same():
+    x = np.random.random_sample(30 * 30 * 3).reshape([1, 30, 30, 3])
+    x1 = KC.variable(x)
+    loss = objectives.DSSIMObjective()
+    assert_allclose([0.0], KC.eval(loss(x1, x1)), atol=1.0e-4)
+
+
+def test_dssim_opposite():
+    x = np.zeros([1, 30, 30, 3])
+    x1 = KC.variable(x)
+    y = np.ones([1, 30, 30, 3])
+    y1 = KC.variable(y)
+    loss = objectives.DSSIMObjective()
+    assert_allclose([0.5], KC.eval(loss(x1, y1)), atol=1.0e-4)
+
+
+def test_dssim_compile():
+    from keras.models import Sequential
+    from keras.layers import Convolution2D
+    x = np.zeros([1, 30, 30, 3])
+    loss = objectives.DSSIMObjective()
+    model = Sequential()
+    model.add(Convolution2D(3, 3, 3, border_mode="same", input_shape=(30, 30, 3)))
+    model.compile("rmsprop", loss)
+    model.fit([x], [x], 1, 1)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])