New segmentation metric: Hausdorff Distance

CHANGELOG.md

@@ -30,6 +30,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `truncation` argument to `BERTScore` ([#2776](https://github.com/Lightning-AI/torchmetrics/pull/2776))
 
 
+- Added `HausdorffDistance` to segmentation package ([#2122](https://github.com/Lightning-AI/torchmetrics/pull/2122))
+
+
 ### Changed
 
 - Tracker higher is better integration ([#2649](https://github.com/Lightning-AI/torchmetrics/pull/2649))

docs/source/links.rst

@@ -172,4 +172,6 @@
 .. _averaging curve objects: https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html
 .. _SCC: https://www.ingentaconnect.com/content/tandf/tres/1998/00000019/00000004/art00013
 .. _Generalized Dice Score: https://arxiv.org/abs/1707.03237
+.. _Hausdorff Distance: https://en.wikipedia.org/wiki/Hausdorff_distance
+.. _averaging curve objects: https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html
 .. _Procrustes Disparity: https://en.wikipedia.org/wiki/Procrustes_analysis

docs/source/segmentation/hausdorff_distance.rst

@@ -0,0 +1,21 @@
+.. customcarditem::
+   :header: Hausdorff Distance
+   :image: https://pl-flash-data.s3.amazonaws.com/assets/thumbnails/text_classification.svg
+   :tags: segmentation
+
+.. include:: ../links.rst
+
+##################
+Hausdorff Distance
+##################
+
+Module Interface
+________________
+
+.. autoclass:: torchmetrics.segmentation.HausdorffDistance
+    :exclude-members: update, compute
+
+Functional Interface
+____________________
+
+.. autofunction:: torchmetrics.functional.segmentation.hausdorff_distance

src/torchmetrics/functional/segmentation/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from torchmetrics.functional.segmentation.generalized_dice import generalized_dice_score
+from torchmetrics.functional.segmentation.hausdorff_distance import hausdorff_distance
 from torchmetrics.functional.segmentation.mean_iou import mean_iou
 
-__all__ = ["generalized_dice_score", "mean_iou"]
+__all__ = ["generalized_dice_score", "mean_iou", "hausdorff_distance"]

src/torchmetrics/functional/segmentation/hausdorff_distance.py

@@ -0,0 +1,128 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Literal, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from torchmetrics.functional.segmentation.utils import check_if_binarized, surface_distance
+from torchmetrics.utilities.checks import _check_same_shape
+
+
+def _hausdorff_distance_validate_args(
+    distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
+    spacing: Optional[Union[Tensor, list[float]]] = None,
+) -> None:
+    """Validate the arguments of `hausdorff_distance` function."""
+    if distance_metric not in ["euclidean", "chessboard", "taxicab"]:
+        raise ValueError(
+            f"Arg `distance_metric` must be one of 'euclidean', 'chessboard', 'taxicab', but got {distance_metric}."
+        )
+    if spacing is not None and not isinstance(spacing, (list, Tensor)):
+        raise ValueError(f"Arg `spacing` must be a list or tensor, but got {type(spacing)}.")
+
+
+def _hausdorff_distance_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
+    """Update and returns variables required to compute `Hausdorff Distance`_.
+
+    Args:
+        preds: predicted binarized segmentation map
+        target: target binarized segmentation map
+
+    Returns:
+        preds: predicted binarized segmentation map
+        target: target binarized segmentation map
+
+    """
+    check_if_binarized(preds)
+    check_if_binarized(target)
+    _check_same_shape(preds, target)
+    return preds, target
+
+
+def _hausdorff_distance_compute(
+    preds: Tensor,
+    target: Tensor,
+    distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
+    spacing: Optional[Union[Tensor, list[float]]] = None,
+) -> Tensor:
+    """Compute `Hausdorff Distance`_.
+
+    Args:
+        preds: predicted binarized segmentation map
+        target: target binarized segmentation map
+        distance_metric: distance metric to calculate surface distance. One of `["euclidean", "chessboard", "taxicab"]`.
+        spacing: spacing between pixels along each spatial dimension
+
+    Returns:
+        Hausdorff distance
+
+    Example:
+        >>> preds = torch.tensor([[1, 1, 1, 1, 1],
+        ...                       [1, 0, 0, 0, 1],
+        ...                       [1, 0, 0, 0, 1],
+        ...                       [1, 0, 0, 0, 1],
+        ...                       [1, 1, 1, 1, 1]], dtype=torch.bool)
+        >>> target = torch.tensor([[1, 1, 1, 1, 0],
+        ...                        [1, 0, 0, 1, 0],
+        ...                        [1, 0, 0, 1, 0],
+        ...                        [1, 0, 0, 1, 0],
+        ...                        [1, 1, 1, 1, 0]], dtype=torch.bool)
+        >>> hausdorff_distance(preds, target, distance_metric="euclidean")
+        tensor(1.)
+
+    """
+    fwd = surface_distance(preds, target, distance_metric=distance_metric, spacing=spacing)
+    bwd = surface_distance(target, preds, distance_metric=distance_metric, spacing=spacing)
+    return torch.max(torch.tensor([fwd.max(), bwd.max()]))
+
+
+def hausdorff_distance(
+    preds: Tensor,
+    target: Tensor,
+    distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
+    spacing: Optional[Union[Tensor, list[float]]] = None,
+) -> Tensor:
+    """Calculate `Hausdorff Distance`_.
+
+    Args:
+        preds: predicted binarized segmentation map
+        target: target binarized segmentation map
+        distance_metric: distance metric to calculate surface distance. One of `["euclidean", "chessboard", "taxicab"]`.
+        spacing: spacing between pixels along each spatial dimension
+
+    Returns:
+        Hausdorff Distance
+
+    Example:
+        >>> import torch
+        >>> from torchmetrics.functional.segmentation import hausdorff_distance
+        >>> preds = torch.tensor([[1, 1, 1, 1, 1],
+        ...                       [1, 0, 0, 0, 1],
+        ...                       [1, 0, 0, 0, 1],
+        ...                       [1, 0, 0, 0, 1],
+        ...                       [1, 1, 1, 1, 1]], dtype=torch.bool)
+        >>> target = torch.tensor([[1, 1, 1, 1, 0],
+        ...                        [1, 0, 0, 1, 0],
+        ...                        [1, 0, 0, 1, 0],
+        ...                        [1, 0, 0, 1, 0],
+        ...                        [1, 1, 1, 1, 0]], dtype=torch.bool)
+        >>> hausdorff_distance(preds, target, distance_metric="euclidean")
+        tensor(1.)
+
+    """
+    _hausdorff_distance_validate_args(distance_metric, spacing)
+    preds, target = _hausdorff_distance_update(preds, target)
+    return _hausdorff_distance_compute(preds, target, distance_metric=distance_metric, spacing=spacing)

src/torchmetrics/functional/segmentation/utils.py

@@ -32,7 +32,7 @@ def _ignore_background(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
 
 
 def check_if_binarized(x: Tensor) -> None:
-    """Check if the input is binarized.
+    """Check if tensor is binarized.
 
     Example:
         >>> from torchmetrics.functional.segmentation.utils import check_if_binarized
@@ -249,25 +249,25 @@ def distance_transform(
             raise ValueError(f"Expected argument `sampling` to have length 2 but got length `{len(sampling)}`.")
 
     if engine == "pytorch":
+        x = x.float()
         # calculate distance from every foreground pixel to every background pixel
         i0, j0 = torch.where(x == 0)
         i1, j1 = torch.where(x == 1)
-        dis_row = (i1.unsqueeze(1) - i0.unsqueeze(0)).abs_().mul_(sampling[0])
-        dis_col = (j1.unsqueeze(1) - j0.unsqueeze(0)).abs_().mul_(sampling[1])
+        dis_row = (i1.view(-1, 1) - i0.view(1, -1)).abs()
+        dis_col = (j1.view(-1, 1) - j0.view(1, -1)).abs()
 
         # # calculate distance
         h, _ = x.shape
         if metric == "euclidean":
-            dis_row = dis_row.float()
-            dis_row.pow_(2).add_(dis_col.pow_(2)).sqrt_()
+            dis = ((sampling[0] * dis_row) ** 2 + (sampling[1] * dis_col) ** 2).sqrt()
         if metric == "chessboard":
-            dis_row = dis_row.max(dis_col)
+            dis = torch.max(sampling[0] * dis_row, sampling[1] * dis_col).float()
         if metric == "taxicab":
-            dis_row.add_(dis_col)
+            dis = (sampling[0] * dis_row + sampling[1] * dis_col).float()
 
         # select only the closest distance
-        mindis, _ = torch.min(dis_row, dim=1)
-        z = torch.zeros_like(x, dtype=mindis.dtype).view(-1)
+        mindis, _ = torch.min(dis, dim=1)
+        z = torch.zeros_like(x).view(-1)
         z[i1 * h + j1] = mindis
         return z.view(x.shape)
 
@@ -279,7 +279,7 @@ def distance_transform(
 
     if metric == "euclidean":
         return ndimage.distance_transform_edt(x.cpu().numpy(), sampling)
-    return ndimage.distance_transform_cdt(x.cpu().numpy(), metric=metric)
+    return ndimage.distance_transform_cdt(x.cpu().numpy(), sampling, metric=metric)
 
 
 def mask_edges(
@@ -345,7 +345,7 @@ def surface_distance(
     target: Tensor,
     distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
     spacing: Optional[Union[Tensor, List[float]]] = None,
-) -> Tensor:
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
     """Calculate the surface distance between two binary edge masks.
 
     May return infinity if the predicted mask is empty and the target mask is not, or vice versa.
@@ -390,6 +390,38 @@ def surface_distance(
     return dis[preds]
 
 
+def edge_surface_distance(
+    preds: Tensor,
+    target: Tensor,
+    distance_metric: Literal["euclidean", "chessboard", "taxicab"] = "euclidean",
+    spacing: Optional[Union[Tensor, List[float]]] = None,
+    symmetric: bool = False,
+) -> Tensor:
+    """Extracts the edges from the input masks and calculates the surface distance between them.
+
+    Args:
+        preds: The predicted binary edge mask.
+        target: The target binary edge mask.
+        distance_metric: The distance metric to use. One of `["euclidean", "chessboard", "taxicab"]`.
+        spacing: The spacing between pixels along each spatial dimension.
+        symmetric: Whether to calculate the symmetric distance between the edges.
+
+    Returns:
+        A tensor with length equal to the number of edges in predictions e.g. `preds.sum()`. Each element is the
+        distance from the corresponding edge in `preds` to the closest edge in `target`. If `symmetric` is `True`, the
+        function returns a tuple containing the distances from the predicted edges to the target edges and vice versa.
+
+    """
+    output = mask_edges(preds, target)
+    edges_preds, edges_target = output[0].bool(), output[1].bool()
+    if symmetric:
+        return (
+            surface_distance(edges_preds, edges_target, distance_metric=distance_metric, spacing=spacing),
+            surface_distance(edges_target, edges_preds, distance_metric=distance_metric, spacing=spacing),
+        )
+    return surface_distance(edges_preds, edges_target, distance_metric=distance_metric, spacing=spacing)
+
+
 @functools.lru_cache
 def get_neighbour_tables(
     spacing: Union[Tuple[int, int], Tuple[int, int, int]], device: Optional[torch.device] = None

src/torchmetrics/segmentation/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from torchmetrics.segmentation.generalized_dice import GeneralizedDiceScore
+from torchmetrics.segmentation.hausdorff_distance import HausdorffDistance
 from torchmetrics.segmentation.mean_iou import MeanIoU
 
-__all__ = ["GeneralizedDiceScore", "MeanIoU"]
+__all__ = ["GeneralizedDiceScore", "MeanIoU", "HausdorffDistance"]