generalised AHD for both 2d and 3d images

miseval/hausdorff.py

@@ -53,28 +53,42 @@ def calc_SimpleHausdorffDistance(truth, pred, c=1, **kwargs):
 
 Implementation: https://github.com/Issam28/Brain-tumor-segmentation
 """
-def border_map(binary_img,neigh):
+def border_map(binary_img):
     """
-    Creates the border for a 3D image
+    Creates the border for a 3D or 2D image
     """
+    ndims = binary_img.ndim
     binary_map = np.asarray(binary_img, dtype=np.uint8)
-    neigh = neigh
-    west = ndimage.shift(binary_map, [-1, 0], order=0)
-    east = ndimage.shift(binary_map, [1, 0], order=0)
-    north = ndimage.shift(binary_map, [0, 1], order=0)
-    south = ndimage.shift(binary_map, [0, -1], order=0)
-    cumulative = west + east + north + south
-    border = ((cumulative < 4) * binary_map) == 1
+
+    if ndims == 2:
+        left = ndimage.shift(binary_map, [-1, 0], order=0)
+        right = ndimage.shift(binary_map, [1, 0], order=0)
+        superior = ndimage.shift(binary_map, [0, 1], order=0)
+        inferior = ndimage.shift(binary_map, [0, -1], order=0)
+        cumulative = left + right + superior + inferior
+        ndir = 4
+    elif ndims == 3:
+        left = ndimage.shift(binary_map, [-1, 0, 0], order=0)
+        right = ndimage.shift(binary_map, [1, 0, 0], order=0)
+        anterior = ndimage.shift(binary_map, [0, 1, 0], order=0)
+        posterior = ndimage.shift(binary_map, [0, -1, 0], order=0)
+        superior = ndimage.shift(binary_map, [0, 0, 1], order=0)
+        inferior = ndimage.shift(binary_map, [0, 0, -1], order=0)
+        cumulative = left + right + anterior + posterior + superior + inferior
+        ndir = 6
+    else:
+        raise RuntimeError(f'Image must be of 2 or 3 dimensions, got {ndims}')
+
+    border = ((cumulative < ndir) * binary_map) == 1
     return border
 
 def border_distance(ref,seg):
     """
     This functions determines the map of distance from the borders of the
     segmentation and the reference and the border maps themselves
     """
-    neigh=8
-    border_ref = border_map(ref,neigh)
-    border_seg = border_map(seg,neigh)
+    border_ref = border_map(ref)
+    border_seg = border_map(seg)
     oppose_ref = 1 - ref
     oppose_seg = 1 - seg
     # euclidean distance transform

tests/test_hausdorff.py

@@ -39,6 +39,15 @@ def setUpClass(self):
         np.random.seed(2)
         self.pd_bi = np.random.randint(2, size=(32,32))
         self.pd_mc = np.random.randint(5, size=(32,32))
+        #same for 3d
+        # Create ground truth
+        np.random.seed(1)
+        self.gt_bi_3d = np.random.randint(2, size=(32,32,32))
+        self.gt_mc_3d = np.random.randint(5, size=(32,32,32))
+        # Create prediction
+        np.random.seed(2)
+        self.pd_bi_3d = np.random.randint(2, size=(32,32,32))
+        self.pd_mc_3d = np.random.randint(5, size=(32,32,32))
 
     #-------------------------------------------------#
     #      Calculate : Simple Hausdorff Distance      #
@@ -73,3 +82,15 @@ def test_calc_AverageHausdorffDistance(self):
         self.assertTrue(callable(metric_dict["AHD"]))
         self.assertTrue("AverageHausdorffDistance" in metric_dict)
         self.assertTrue(callable(metric_dict["AverageHausdorffDistance"]))
+
+    #-------------------------------------------------#
+    #     Calculate : Average Hausdorff Distance 3d   #
+    #-------------------------------------------------#
+    def test_calc_AverageHausdorffDistance_3d(self):
+        # Check binary score
+        score_bi = calc_AverageHausdorffDistance(self.gt_bi_3d, self.pd_bi_3d, c=1)
+        self.assertTrue(isinstance(score_bi, np.float64))
+        # Check multi-class score
+        for i in range(5):
+            score_mc = calc_AverageHausdorffDistance(self.gt_mc_3d, self.pd_mc_3d, c=i)
+            self.assertTrue(isinstance(score_mc, np.float64))