Made HausdorffDistance independent of the package file constants

Also renamed it to HaussdorfDistanceErosion, as it is an erosion based approximation (and uploaded the affected files)
Added description
Removed useless imports

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)