Fxed Hausdorff implementation. The successive erosions were not applied to the eroded diff variable, but to the original diff variable (hence, no successive erosion was done at all)

Implemented 3D Structural Similarity Metric Index, based on https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py
The patch size, overlap and dynamic range can be modified at will.

DeepDeformationMapRegistration/losses.py

@@ -34,7 +34,7 @@ class HausdorffDistanceErosion:
 
     def _erosion_distance_single(self, y_true, y_pred):
         diff = tf.math.pow(y_pred - y_true, 2)
-        alpha = 2.
+        alpha = 2
 
         norm = 1 / (self.ndims * 2 + 1)
         kernel = generate_binary_structure(self.ndims, 1).astype(int) * norm
@@ -42,10 +42,12 @@ class HausdorffDistanceErosion:
         kernel = tf.expand_dims(tf.expand_dims(kernel, -1), -1)   # [H, W, D, C_in, C_out]
 
         ret = 0.
-        for i in range(self.nerosions):
-            for j in range(i + 1):
-                er = self._erode_per_channel(diff, kernel)
-            ret += tf.reduce_sum(tf.multiply(er, tf.pow(i + 1., alpha)), self.sum_range)
+        for k in range(1, self.nerosions+1):
+            er = diff
+            # k successive erosions
+            for j in range(k):
+                er = self._erode_per_channel(er, kernel)
+            ret += tf.reduce_sum(tf.multiply(er, tf.cast(tf.pow(k, alpha), tf.float32)), self.sum_range)
 
         img_vol = tf.cast(tf.reduce_prod(tf.shape(y_true)[:-1]), tf.float32)  # Volume of each channel
         return tf.divide(ret, img_vol)  # Divide by the image size
@@ -54,7 +56,7 @@ class HausdorffDistanceErosion:
         batched_dist = tf.map_fn(lambda x: self._erosion_distance_single(x[0], x[1]), (y_true, y_pred),
                                  dtype=tf.float32)
 
-        return batched_dist
+        return tf.reduce_mean(batched_dist)
 
 
 class NCC:
@@ -77,4 +79,60 @@ class NCC:
         return tf.math.divide_no_nan(numerator, denominator)
 
     def loss(self, y_true, y_pred):
-        return tf.map_fn(lambda x: 1 - self.ncc(x[0], x[1]), (y_true, y_pred), tf.float32)
+        # According to the documentation, the loss returns a scalar
+        # Ref: https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile
+        return tf.reduce_mean(tf.map_fn(lambda x: 1 - self.ncc(x[0], x[1]), (y_true, y_pred), tf.float32))
+
+
+class StructuralSimilarity:
+    # Based on https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py
+    def __init__(self, k1=0.01, k2=0.03, patch_size=3, dynamic_range=1., overlap=0.0):
+        """
+        Structural (Di)Similarity Index Measure:
+
+        :param k1: Internal parameter. Defaults to 0.01
+        :param k2: Internal parameter. Defaults to 0.02
+        :param patch_size: Size of the extracted patches
+        :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        """
+        self.__c1 = (k1 * dynamic_range) ** 2
+        self.__c2 = (k2 * dynamic_range) ** 2
+        self.__kernel_shape = [1] + [patch_size] * 3 + [1]
+        stride = int(patch_size * (1 - overlap))
+        self.__stride = [1] + [stride if stride else 1] * 3 + [1]
+        self.__max_val = dynamic_range
+
+    def __int_shape(self, x):
+        return tf.keras.backend.int_shape(x) if tf.keras.backend.backend() == 'tensorflow' else tf.keras.backend.shape(x)
+
+    def ssim(self, y_true, y_pred):
+
+        patches_true = tf.extract_volume_patches(y_true, self.__kernel_shape, self.__stride, 'VALID',
+                                                 'patches_true')
+        patches_pred = tf.extract_volume_patches(y_pred, self.__kernel_shape, self.__stride, 'VALID',
+                                                 'patches_pred')
+
+        #bs, w, h, d, *c = self.__int_shape(patches_pred)
+        #patches_true = tf.reshape(patches_true, [-1, w, h, d, tf.reduce_prod(c)])
+        #patches_pred = tf.reshape(patches_pred, [-1, w, h, d, tf.reduce_prod(c)])
+
+        # Mean
+        u_true = tf.reduce_mean(patches_true, axis=-1)
+        u_pred = tf.reduce_mean(patches_pred, axis=-1)
+
+        # Variance
+        v_true = tf.math.reduce_variance(patches_true, axis=-1)
+        v_pred = tf.math.reduce_variance(patches_pred, axis=-1)
+
+        # Covariance
+        covar = tf.reduce_mean(patches_true * patches_pred, axis=-1) - u_true * u_pred
+
+        # SSIM
+        numerator = (2 * u_true * u_pred + self.__c1) * (2 * covar + self.__c2)
+        denominator = ((tf.square(u_true) + tf.square(u_pred) + self.__c1) * (v_pred + v_true + self.__c2))
+        ssim = numerator / denominator
+
+        return tf.reduce_mean(ssim)
+
+    def dssim(self, y_true, y_pred):
+        return tf.reduce_mean((1. - self.ssim(y_true, y_pred)) / 2.0)