Repository: mlcommons/GaNDLF

Issue:
Add "normalized surface dice" (NSD) as a segmentation metric
**Is your feature request related to a problem? Please describe.**
Currently, the segmentation-only metrics that can be used are dice and hausdorff. This is usually fine, but NSD might provide additional context.

**Describe the solution you'd like**
Add this implementation to metric calculation.

**Describe alternatives you've considered**
N.A.

**Additional context**
- the MONAI reference implementation [[ref](https://docs.monai.io/en/stable/metrics.html#surface-dice)] works only on 2D images
- nnUNet's implementation [[ref](https://github.com/MIC-DKFZ/nnUNet/blob/master/nnunet/evaluation/surface_dice.py)] works for both 2D and 3D


Discussion:

Resolving patch:
diff --git a/GANDLF/metrics/__init__.py b/GANDLF/metrics/__init__.py
--- a/GANDLF/metrics/__init__.py
+++ b/GANDLF/metrics/__init__.py
@@ -9,6 +9,8 @@
     hd100_per_label,
     hd95,
     hd95_per_label,
+    nsd,
+    nsd_per_label,
 )
 from .regression import classification_accuracy, balanced_acc_score, per_label_accuracy
 from .generic import recall_score, precision_score, iou_score, f1_score, accuracy
@@ -30,6 +32,12 @@
     "hd100_per_label": hd100_per_label,
     "hausdorff": hd100,
     "hausdorff100": hd100,
+    "nsd": nsd,
+    "nsd_per_label": nsd_per_label,
+    "normalized_surface_dice": nsd,
+    "normalized_surface_dice_per_label": nsd_per_label,
+    "normalized_sd": nsd,
+    "normalized_sd_per_label": nsd_per_label,
     "cel": CEL,
     "f1_score": f1_score,
     "f1": f1_score,
@@ -42,6 +50,25 @@
 }
 
 
+# global define for the metrics that use surface distances, and hence require "connectivity" to be defined
+surface_distance_ids = [
+    "hd95",
+    "hd95_per_label",
+    "hausdorff95_per_label",
+    "hausdorff95",
+    "hd100",
+    "hd100_per_label",
+    "hausdorff",
+    "hausdorff100",
+    "nsd",
+    "nsd_per_label",
+    "normalized_surface_dice",
+    "normalized_surface_dice_per_label",
+    "normalized_sd",
+    "normalized_sd_per_label",
+]
+
+
 def overall_stats(predictions, ground_truth, params):
     """
     Generates a dictionary of metrics calculated on the overall predictions and ground truths.
diff --git a/GANDLF/metrics/segmentation.py b/GANDLF/metrics/segmentation.py
--- a/GANDLF/metrics/segmentation.py
+++ b/GANDLF/metrics/segmentation.py
@@ -1,7 +1,9 @@
 """
 All the segmentation metrics are to be called from here
 """
-import torch, numpy
+import sys
+import torch
+import numpy as np
 from GANDLF.losses.segmentation import dice
 from scipy.ndimage import _ni_support
 from scipy.ndimage.morphology import (
@@ -11,28 +13,37 @@
 )
 
 
+def _convert_tensor_to_int_label_array(input_tensor):
+    """
+    This function converts a tensor of labels to a numpy array of labels.
+
+    Args:
+        input_tensor (torch.Tensor): Input data containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+
+    Returns:
+        numpy.ndarray: The numpy array of labels.
+    """
+    result_array = input_tensor.detach().cpu().numpy()
+    if result_array.shape[-1] == 1:
+        result_array = result_array.squeeze(-1)
+    # ensure that we are dealing with a binary array
+    result_array[result_array < 0.5] = 0
+    result_array[result_array >= 0.5] = 1
+    return result_array.astype(np.int64)
+
+
 def multi_class_dice(output, label, params, per_label=False):
     """
-    This function computes a multi-class dice
-
-    Parameters
-    ----------
-    output : torch.Tensor
-        Input data containing objects. Can be any type but will be converted
-        into binary: background where 0, object everywhere else.
-    label : torch.Tensor
-        Input data containing objects. Can be any type but will be converted
-        into binary: background where 0, object everywhere else.
-    params : dict
-        The parameter dictionary containing training and data information.
-    per_label : bool
-        Whether the dice needs to be calculated per label or not.
-
-    Returns
-    -------
-    total_dice : float or list
-        The average dice for all labels or a list of per-label dice scores.
+    This function computes a multi-class dice.
+
+    Args:
+        output (torch.Tensor): Input data containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        label (torch.Tensor): Input data containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        params (dict): The parameter dictionary containing training and data information.
+        per_label (bool, optional): Whether the dice needs to be calculated per label or not. Defaults to False.
 
+    Returns:
+        float or list: The average dice for all labels or a list of per-label dice scores.
     """
     total_dice = 0
     avg_counter = 0
@@ -55,23 +66,15 @@ def multi_class_dice(output, label, params, per_label=False):
 
 def multi_class_dice_per_label(output, label, params):
     """
-    This function computes a multi-class dice
-
-    Parameters
-    ----------
-    output : torch.Tensor
-        Input data containing objects. Can be any type but will be converted
-        into binary: background where 0, object everywhere else.
-    label : torch.Tensor
-        Input data containing objects. Can be any type but will be converted
-        into binary: background where 0, object everywhere else.
-    params : dict
-        The parameter dictionary containing training and data information.
-
-    Returns
-    -------
-    total_dice : list
-        The list of per-label dice scores.
+    This function computes a multi-class dice.
+
+    Args:
+        output (torch.Tensor): Input data containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        label (torch.Tensor): Input data containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        params (dict): The parameter dictionary containing training and data information.
+
+    Returns:
+        list: The list of per-label dice scores.
     """
     return multi_class_dice(output, label, params, per_label=True)
 
@@ -79,15 +82,22 @@ def multi_class_dice_per_label(output, label, params):
 def __surface_distances(result, reference, voxelspacing=None, connectivity=1):
     """
     The distances between the surface voxel of binary objects in result and their
-    nearest partner surface voxel of a binary object in reference.
+    nearest partner surface voxel of a binary object in reference. Adapted from https://github.com/loli/medpy/blob/39131b94f0ab5328ab14a874229320efc2f74d98/medpy/metric/binary.py#L1195.
 
-    Adopted from https://github.com/loli/medpy/blob/39131b94f0ab5328ab14a874229320efc2f74d98/medpy/metric/binary.py#L1195
+    Args:
+        result (torch.Tensor): Input prediction containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        reference (torch.Tensor): Input ground truth containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        voxelspacing (tuple): The size of each voxel, defaults to isotropic spacing of 1mm.
+        connectivity (int): The connectivity of regions. See scipy.ndimage.generate_binary_structure for more information.
+
+    Returns:
+        float: The symmetric Hausdorff Distance between the object(s) in ```result``` and the object(s) in ```reference```. The distance unit is the same as for the spacing of elements along each dimension, which is usually given in mm.
     """
-    result = numpy.atleast_1d(result.astype(numpy.bool))
-    reference = numpy.atleast_1d(reference.astype(numpy.bool))
+    result = np.atleast_1d(result.astype(np.bool))
+    reference = np.atleast_1d(reference.astype(np.bool))
     if voxelspacing is not None:
         voxelspacing = _ni_support._normalize_sequence(voxelspacing, result.ndim)
-        voxelspacing = numpy.asarray(voxelspacing, dtype=numpy.float64)
+        voxelspacing = np.asarray(voxelspacing, dtype=np.float64)
         if not voxelspacing.flags.contiguous:
             voxelspacing = voxelspacing.copy()
 
@@ -95,17 +105,9 @@ def __surface_distances(result, reference, voxelspacing=None, connectivity=1):
     footprint = generate_binary_structure(result.ndim, connectivity)
 
     # test for emptiness
-    if 0 == numpy.count_nonzero(result):
-        # print(
-        #     "The first supplied array does not contain any binary object.",
-        #     file=sys.stderr,
-        # )
+    if 0 == np.count_nonzero(result):
         return 0
-    if 0 == numpy.count_nonzero(reference):
-        # print(
-        #     "The second supplied array does not contain any binary object.",
-        #     file=sys.stderr,
-        # )
+    if 0 == np.count_nonzero(reference):
         return 0
 
     # extract only 1-pixel border line of objects
@@ -123,48 +125,61 @@ def __surface_distances(result, reference, voxelspacing=None, connectivity=1):
     return sds
 
 
-def hd_generic(inp, target, params, percentile=95, per_label=False):
+def _nsd_base(a_to_b, b_to_a, threshold):
     """
-    Generic Hausdorff Distance calculation
-    Computes the Hausdorff Distance (HD) between the binary objects in two
-    images. Compared to the Hausdorff Distance, this metric is slightly more stable to small outliers and is
-    commonly used in Biomedical Segmentation challenges.
-    Parameters
-    ----------
-    result : torch.tensor
-        Input data containing objects. Can be any type but will be converted
-        into binary: background where 0, object everywhere else.
-    reference : torch.tensor
-        Input data containing objects. Can be any type but will be converted
-        into binary: background where 0, object everywhere else.
-    params : dict
-        The parameter dictionary containing training and data information.
-    percentile : int
-        The percentile of surface distances to include during Hausdorff calculation.
-    per_label : bool
-        Whether the hausdorff needs to be calculated per label or not.
-    Returns
-    -------
-    hd : float
-        The symmetric Hausdorff Distance between the object(s) in ```result``` and the
-        object(s) in ```reference```. The distance unit is the same as for the spacing of
-        elements along each dimension, which is usually given in mm.
-    See also
-    --------
-    :func:`hd`
-    Notes
-    -----
-    This is a real metric. The binary images can therefore be supplied in any order.
+    This implementation differs from the official surface dice implementation! These two are not comparable!!!!!
+    The normalized surface dice is symmetric, so it should not matter whether a or b is the reference image
+    This implementation natively supports 2D and 3D images.
+
+    Args:
+        a_to_b (np.ndarray): Surface distances from a to b
+        b_to_a (np.ndarray): Surface distances from b to a
+        threshold (float): distances below this threshold will be counted as true positives. Threshold is in mm, not voxels!
+
+    Returns:
+        float: the normalized surface dice between a and b
     """
-    result_array = inp.detach().cpu().numpy()
-    target_array = target.detach().cpu().numpy()
-    if result_array.shape[-1] == 1:
-        result_array = result_array.squeeze(-1)
-    if target_array.shape[-1] == 1:
-        target_array = target_array.squeeze(-1)
-    # ensure that we are dealing with a binary array
-    result_array[result_array < 0.5] = 0
-    result_array[result_array >= 0.5] = 1
+    if isinstance(a_to_b, int):
+        return 0
+    if isinstance(b_to_a, int):
+        return 0
+    numel_a = len(a_to_b)
+    numel_b = len(b_to_a)
+    tp_a = np.sum(a_to_b <= threshold) / numel_a
+    tp_b = np.sum(b_to_a <= threshold) / numel_b
+    fp = np.sum(a_to_b > threshold) / numel_a
+    fn = np.sum(b_to_a > threshold) / numel_b
+    dc = (tp_a + tp_b) / (tp_a + tp_b + fp + fn + sys.float_info.min)
+    return dc
+
+
+def _calculator_generic(
+    inp,
+    target,
+    params,
+    percentile=95,
+    surface_dice=False,
+    per_label=False,
+):
+    """
+    Generic Surface Dice (SD)/Hausdorff (HD) Distance calculation from 2 tensors. Compared to the standard Hausdorff Distance, this metric is slightly more stable to small outliers and is commonly used in Biomedical Segmentation challenges.
+
+    Args:
+        inp (torch.Tensor): Input prediction containing objects. Can be any type but will be converted into binary: background where 0, object everywhere else.
+        target (torch.Tensor): Input ground truth containing objects. Can be any type but will be converted into binary: binary: background where 0, object everywhere else.
+        params (dict): The parameter dictionary containing training and data information.
+        percentile (int, optional): The percentile of surface distances to include during HD calculation. Defaults to 95.
+        surface_dice (bool, optional): Whether the SD needs to be calculated or not. Defaults to False.
+        per_label (bool, optional): Whether the hausdorff needs to be calculated per label or not. Defaults to False.
+
+    Returns:
+        float or list: The symmetric Hausdorff Distance between the object(s) in ```result``` and the object(s) in ```reference```. The distance unit is the same as for the spacing of elements along each dimension, which is usually given in mm.
+
+    See also:
+        :func:`hd`
+    """
+    result_array = _convert_tensor_to_int_label_array(inp)
+    target_array = _convert_tensor_to_int_label_array(target)
 
     hd = 0
     avg_counter = 0
@@ -175,13 +190,19 @@ def hd_generic(inp, target, params, percentile=95, per_label=False):
                 hd1 = __surface_distances(
                     result_array[b, i, ...],
                     target_array[b, i, ...],
+                    params["subject_spacing"][b],
                 )
                 hd2 = __surface_distances(
                     target_array[b, i, ...],
                     result_array[b, i, ...],
                     params["subject_spacing"][b],
                 )
-                current_hd = numpy.percentile(numpy.hstack((hd1, hd2)), percentile)
+                if surface_dice:
+                    # ensure threshold always at least 1
+                    threshold = max(min(params["subject_spacing"][0]), 1).item()
+                    current_hd = _nsd_base(hd1, hd2, threshold)
+                else:
+                    current_hd = np.percentile(np.hstack((hd1, hd2)), percentile)
                 hd += current_hd
                 hd_per_label.append(current_hd)
                 avg_counter += 1
@@ -193,16 +214,26 @@ def hd_generic(inp, target, params, percentile=95, per_label=False):
 
 
 def hd95(inp, target, params):
-    return hd_generic(inp, target, params, 95)
+    return _calculator_generic(inp, target, params, percentile=95)
 
 
 def hd95_per_label(inp, target, params):
-    return hd_generic(inp, target, params, 95, True)
+    return _calculator_generic(inp, target, params, percentile=95, per_label=True)
 
 
 def hd100(inp, target, params):
-    return hd_generic(inp, target, params, 100)
+    return _calculator_generic(inp, target, params, percentile=100)
 
 
 def hd100_per_label(inp, target, params):
-    return hd_generic(inp, target, params, 100, True)
+    return _calculator_generic(inp, target, params, percentile=100, per_label=True)
+
+
+def nsd(inp, target, params):
+    return _calculator_generic(inp, target, params, percentile=100, surface_dice=True)
+
+
+def nsd_per_label(inp, target, params):
+    return _calculator_generic(
+        inp, target, params, percentile=100, per_label=True, surface_dice=True
+    )
diff --git a/GANDLF/parseConfig.py b/GANDLF/parseConfig.py
--- a/GANDLF/parseConfig.py
+++ b/GANDLF/parseConfig.py
@@ -5,6 +5,8 @@
 from .utils import version_check
 from GANDLF.data.post_process import postprocessing_after_reverse_one_hot_encoding
 
+from GANDLF.metrics import surface_distance_ids
+
 ## dictionary to define defaults for appropriate options, which are evaluated
 parameter_defaults = {
     "weighted_loss": False,  # whether weighted loss is to be used or not
@@ -262,6 +264,13 @@ def parseConfig(config_file_path, version_check_flag=True):
                     temp_dict["iou"], "reduction", "elementwise_mean"
                 )
                 temp_dict["iou"] = initialize_key(temp_dict["iou"], "threshold", 0.5)
+            elif comparison_string in surface_distance_ids:
+                temp_dict[comparison_string] = initialize_key(
+                    temp_dict[comparison_string], "connectivity", 1
+                )
+                temp_dict[comparison_string] = initialize_key(
+                    temp_dict[comparison_string], "threshold", None
+                )
 
         params["metrics"] = temp_dict
 


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -1115,7 +1115,12 @@ def test_train_metrics_segmentation_rad_2d(device):
     parameters["model"]["amp"] = True
     parameters["save_output"] = True
     parameters["model"]["num_channels"] = 3
-    parameters["metrics"] = ["dice", "hausdorff", "hausdorff95"]
+    parameters["metrics"] = [
+        "dice",
+        "hausdorff",
+        "hausdorff95",
+        "normalized_surface_dice",
+    ]
     parameters["model"]["architecture"] = "resunet"
     parameters["model"]["onnx_export"] = False
     parameters["model"]["print_summary"] = False
@@ -1763,6 +1768,8 @@ def test_train_checkpointing_segmentation_rad_2d(device):
         "hausdorff95",
         "hd95_per_label",
         "hd100_per_label",
+        "normalized_surface_dice",
+        "normalized_surface_dice_per_label",
     ]
     parameters["model"]["architecture"] = "unet"
     parameters["model"]["onnx_export"] = False
@@ -1809,7 +1816,11 @@ def test_generic_model_patch_divisibility():
     parameters["model"]["amp"] = True
     parameters["model"]["print_summary"] = False
     parameters["model"]["num_channels"] = 3
-    parameters["metrics"] = ["dice", "hausdorff", "hausdorff95"]
+    parameters["metrics"] = [
+        "dice",
+        "hausdorff",
+        "hausdorff95" "normalized_surface_dice_per_label",
+    ]
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
 
     # this assertion should fail
@@ -2168,17 +2179,13 @@ def test_train_inference_classification_histology_large_2d(device):
                 # check in the default outputDir that's created - this is based on a unique timestamp
                 if folder != "output_validation":
                     # if 'predictions.csv' are not found, give error
-                    assert (
-                        os.path.exists(
-                            os.path.join(
-                                output_subject_dir,
-                                str(input_df["SubjectID"][0]),
-                                "predictions.csv",
-                            )
+                    assert os.path.exists(
+                        os.path.join(
+                            output_subject_dir,
+                            str(input_df["SubjectID"][0]),
+                            "predictions.csv",
                         )
-                        is True,
-                        "predictions.csv not found",
-                    )
+                    ), "predictions.csv not found"
     # ensure previous results are removed
     sanitize_outputDir()