Merge pull request #1 from jpdefrutosSINTEF/HausdorffConstantIndependent

DeepDeformationMapRegistration/losses.py

@@ -1,19 +1,19 @@
-import os, sys
-currentdir = os.path.dirname(os.path.realpath(__file__))
-parentdir = os.path.dirname(currentdir)
-sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
-
-PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
-
 import tensorflow as tf
 from scipy.ndimage import generate_binary_structure
 
-import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.utils.operators import soft_threshold
 
 
-class HausdorffDistance:
+class HausdorffDistanceErosion:
     def __init__(self, ndim=3, nerosion=10):
+        """
+        Approximation of the Hausdorff distance based on erosion operations based on the work done by Karimi D., et al.
+            Karimi D., et al., "Reducing the Hausdorff Distance in Medical Image Segmentation with Convolutional Neural
+            Networks". IEEE Transactions on Medical Imaging, 39, 2020. DOI 10.1109/TMI.2019.2930068
+
+        :param ndim: Dimensionality of the images
+        :param nerosion: Number of erosion steps. Defaults to 10.
+        """
         self.ndims = ndim
         self.conv = getattr(tf.nn, 'conv%dd' % self.ndims)
         self.nerosions = nerosion
@@ -37,11 +37,12 @@ class HausdorffDistance:
                 er = self._erode(diff, kernel)
             ret += tf.reduce_sum(tf.multiply(er, tf.pow(i + 1., alpha)))
 
-        return tf.multiply(C.IMG_SIZE ** -self.ndims, ret)  # Divide by the image size
+        img_vol = tf.cast(tf.reduce_prod(y_true.shape), tf.float32)
+        return tf.divide(ret, img_vol)  # Divide by the image size
 
     def loss(self, y_true, y_pred):
         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  # tf.reduce_mean(batched_dist)
+        return batched_dist
 

TrainingScripts/Train_2d.py

@@ -14,7 +14,7 @@ from datetime import datetime
 import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.data_generator import DataGeneratorManager2D
 from DeepDeformationMapRegistration.utils.misc import try_mkdir
-from DeepDeformationMapRegistration.losses import HausdorffDistance
+from DeepDeformationMapRegistration.losses import HausdorffDistanceErosion
 
 
 os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
@@ -52,7 +52,7 @@ vxm_model = vxm.networks.VxmDense(inshape=in_shape, nb_unet_features=nb_features
 
 # Losses and loss weights
 def comb_loss(y_true, y_pred):
-    return 1e-3 * HausdorffDistance(ndim=2, nerosion=2).loss(y_true, y_pred) + vxm.losses.Dice().loss(y_true, y_pred)
+    return 1e-3 * HausdorffDistanceErosion(ndim=2, nerosion=2).loss(y_true, y_pred) + vxm.losses.Dice().loss(y_true, y_pred)
 
 
 losses = [comb_loss, vxm.losses.Grad('l2').loss]

TrainingScripts/Train_2d_uncertaintyWeighting.py

@@ -17,7 +17,7 @@ from datetime import datetime
 import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.data_generator import DataGeneratorManager2D
 from DeepDeformationMapRegistration.utils.misc import try_mkdir
-from DeepDeformationMapRegistration.losses import HausdorffDistance
+from DeepDeformationMapRegistration.losses import HausdorffDistanceErosion
 from DeepDeformationMapRegistration.layers import UncertaintyWeighting
 
 
@@ -66,7 +66,7 @@ def dice_loss(y_true, y_pred):
 #fixed_pred, dm_pred = vxm_model([moving, fixed])
 multiLoss = UncertaintyWeighting(num_loss_fns=2,
                                  num_reg_fns=1,
-                                 loss_fns=[HausdorffDistance(2, 2).loss, dice_loss],
+                                 loss_fns=[HausdorffDistanceErosion(2, 2).loss, dice_loss],
                                  reg_fns=[vxm.losses.Grad('l2').loss],
                                  prior_loss_w=[1., 1.],
                                  prior_reg_w=[0.01],

TrainingScripts/Train_3d_weaklySupervised.py

@@ -17,7 +17,7 @@ import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.data_generator import DataGeneratorManager
 from DeepDeformationMapRegistration.utils.misc import try_mkdir
 from DeepDeformationMapRegistration.networks import WeaklySupervised
-from DeepDeformationMapRegistration.losses import HausdorffDistance
+from DeepDeformationMapRegistration.losses import HausdorffDistanceErosion
 from DeepDeformationMapRegistration.layers import UncertaintyWeighting
 
 
@@ -49,16 +49,16 @@ vxm_model = WeaklySupervised(inshape=in_shape, all_labels=[1], nb_unet_features=
 # Losses and loss weights
 
 grad = tf.keras.Input(shape=(*in_shape, 3), name='multiLoss_grad_input', dtype=tf.float32)
-fix_img = tf.keras.Input(shape=(*in_shape, 1), name='multiLoss_fix_img_input', dtype=tf.float32)
+# fix_img = tf.keras.Input(shape=(*in_shape, 1), name='multiLoss_fix_img_input', dtype=tf.float32)
 def dice_loss(y_true, y_pred):
     # Dice().loss returns -Dice score
     return 1 + vxm.losses.Dice().loss(y_true, y_pred)
 
 multiLoss = UncertaintyWeighting(num_loss_fns=3,
                                  num_reg_fns=1,
-                                 loss_fns=[HausdorffDistance(3, 5).loss, dice_loss, vxm.losses.NCC().loss],
+                                 loss_fns=[HausdorffDistanceErosion(3, 5).loss, dice_loss, vxm.losses.NCC().loss],
                                  reg_fns=[vxm.losses.Grad('l2').loss],
-                                 prior_loss_w=[1., 1., 1.],
+                                 prior_loss_w=[1., 1.],
                                  prior_reg_w=[0.01],
                                  name='MultiLossLayer')
 loss = multiLoss([vxm_model.inputs[1], vxm_model.inputs[1], fix_img,
@@ -66,7 +66,7 @@ loss = multiLoss([vxm_model.inputs[1], vxm_model.inputs[1], fix_img,
                   grad,
                   vxm_model.references.pos_flow])
 
-full_model = tf.keras.Model(inputs=vxm_model.inputs + [fix_img, grad], outputs=vxm_model.outputs + [loss])
+full_model = tf.keras.Model(inputs=vxm_model.inputs + [grad], outputs=vxm_model.outputs + [loss])
 
 # Compile the model
 full_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=None)