def compute_hausdorff_distance(
y_pred: Union[np.ndarray, torch.Tensor],
y: Union[np.ndarray, torch.Tensor],
include_background: bool = False,
distance_metric: str = "euclidean",
percentile: Optional[float] = None,
directed: bool = False,
):
"""
Compute the Hausdorff distance.
Args:
y_pred: input data to compute, typical segmentation model output.
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
should be binarized.
y: ground truth to compute mean the distance. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip distance computation on the first channel of
the predicted output. Defaults to ``False``.
distance_metric: : [``"euclidean"``, ``"chessboard"``, ``"taxicab"``]
the metric used to compute surface distance. Defaults to ``"euclidean"``.
percentile: an optional float number between 0 and 100. If specified, the corresponding
percentile of the Hausdorff Distance rather than the maximum result will be achieved.
Defaults to ``None``.
directed: whether to calculate directed Hausdorff distance. Defaults to ``False``.
"""
if not include_background:
y_pred, y = ignore_background(y_pred=y_pred, y=y)
if isinstance(y, torch.Tensor):
y = y.float()
if isinstance(y_pred, torch.Tensor):
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError("y_pred and y should have same shapes.")
batch_size, n_class = y_pred.shape[:2]
hd = np.empty((batch_size, n_class))
for b, c in np.ndindex(batch_size, n_class):
(edges_pred, edges_gt) = get_mask_edges(y_pred[b, c], y[b, c])
if not np.any(edges_gt):
warnings.warn(
f"the ground truth of class {c} is all 0, this may result in nan/inf distance."
)
if not np.any(edges_pred):
warnings.warn(
f"the prediction of class {c} is all 0, this may result in nan/inf distance."
)
distance_1 = compute_percent_hausdorff_distance(
edges_pred, edges_gt, distance_metric, percentile)
if directed:
hd[b, c] = distance_1
else:
distance_2 = compute_percent_hausdorff_distance(
edges_gt, edges_pred, distance_metric, percentile)
hd[b, c] = max(distance_1, distance_2)
return torch.from_numpy(hd)
def get_confusion_matrix(
y_pred: torch.Tensor,
y: torch.Tensor,
include_background: bool = True,
):
"""
Compute confusion matrix. A tensor with the shape [BC4] will be returned. Where, the third dimension
represents the number of true positive, false positive, true negative and false negative values for
each channel of each sample within the input batch. Where, B equals to the batch size and C equals to
the number of classes that need to be computed.
Args:
y_pred: input data to compute. It must be one-hot format and first dim is batch.
The values should be binarized.
y: ground truth to compute the metric. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip metric computation on the first channel of
the predicted output. Defaults to True.
Raises:
ValueError: when `y_pred` and `y` have different shapes.
"""
if not include_background:
y_pred, y = ignore_background(
y_pred=y_pred,
y=y,
)
y = y.float()
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError("y_pred and y should have same shapes.")
# get confusion matrix related metric
batch_size, n_class = y_pred.shape[:2]
# convert to [BNS], where S is the number of pixels for one sample.
# As for classification tasks, S equals to 1.
y_pred = y_pred.view(batch_size, n_class, -1)
y = y.view(batch_size, n_class, -1)
tp = ((y_pred + y) == 2).float()
tn = ((y_pred + y) == 0).float()
tp = tp.sum(dim=[2])
tn = tn.sum(dim=[2])
p = y.sum(dim=[2])
n = y.shape[-1] - p
fn = p - tp
fp = n - tn
return torch.stack([tp, fp, tn, fn], dim=-1)
def compute_meaniou(y_pred: torch.Tensor,
y: torch.Tensor,
include_background: bool = True,
ignore_empty: bool = True) -> torch.Tensor:
"""Computes IoU score metric from full size Tensor and collects average.
Args:
y_pred: input data to compute, typical segmentation model output.
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
should be binarized.
y: ground truth to compute mean IoU metric. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip IoU computation on the first channel of
the predicted output. Defaults to True.
ignore_empty: whether to ignore empty ground truth cases during calculation.
If `True`, NaN value will be set for empty ground truth cases.
If `False`, 1 will be set if the predictions of empty ground truth cases are also empty.
Returns:
IoU scores per batch and per class, (shape [batch_size, num_classes]).
Raises:
ValueError: when `y_pred` and `y` have different shapes.
"""
if not include_background:
y_pred, y = ignore_background(y_pred=y_pred, y=y)
y = y.float()
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError(
f"y_pred and y should have same shapes, got {y_pred.shape} and {y.shape}."
)
# reducing only spatial dimensions (not batch nor channels)
n_len = len(y_pred.shape)
reduce_axis = list(range(2, n_len))
intersection = torch.sum(y * y_pred, dim=reduce_axis)
y_o = torch.sum(y, reduce_axis)
y_pred_o = torch.sum(y_pred, dim=reduce_axis)
union = y_o + y_pred_o - intersection
if ignore_empty is True:
return torch.where(y_o > 0, (intersection) / union,
torch.tensor(float("nan"), device=y_o.device))
return torch.where(union > 0, (intersection) / union,
torch.tensor(1.0, device=y_o.device))
def compute_meandice(
y_pred: torch.Tensor,
y: torch.Tensor,
include_background: bool = True,
) -> torch.Tensor:
"""Computes Dice score metric from full size Tensor and collects average.
Args:
y_pred: input data to compute, typical segmentation model output.
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
should be binarized.
y: ground truth to compute mean dice metric. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip Dice computation on the first channel of
the predicted output. Defaults to True.
Returns:
Dice scores per batch and per class, (shape [batch_size, n_classes]).
Raises:
ValueError: when `y_pred` and `y` have different shapes.
"""
if not include_background:
y_pred, y = ignore_background(
y_pred=y_pred,
y=y,
)
y = y.float()
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError("y_pred and y should have same shapes.")
# reducing only spatial dimensions (not batch nor channels)
n_len = len(y_pred.shape)
reduce_axis = list(range(2, n_len))
intersection = torch.sum(y * y_pred, dim=reduce_axis)
y_o = torch.sum(y, reduce_axis)
y_pred_o = torch.sum(y_pred, dim=reduce_axis)
denominator = y_o + y_pred_o
f = torch.where(y_o > 0, (2.0 * intersection) / denominator,
torch.tensor(float("nan"), device=y_o.device))
return f # returns array of Dice with shape: [batch, n_classes]
def compute_average_surface_distance(
y_pred: Union[np.ndarray, torch.Tensor],
y: Union[np.ndarray, torch.Tensor],
include_background: bool = False,
symmetric: bool = False,
distance_metric: str = "euclidean",
):
"""
This function is used to compute the Average Surface Distance from `y_pred` to `y`
under the default setting.
In addition, if sets ``symmetric = True``, the average symmetric surface distance between
these two inputs will be returned.
The implementation refers to `DeepMind's implementation <https://github.com/deepmind/surface-distance>`_.
Args:
y_pred: input data to compute, typical segmentation model output.
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
should be binarized.
y: ground truth to compute mean the distance. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip distance computation on the first channel of
the predicted output. Defaults to ``False``.
symmetric: whether to calculate the symmetric average surface distance between
`seg_pred` and `seg_gt`. Defaults to ``False``.
distance_metric: : [``"euclidean"``, ``"chessboard"``, ``"taxicab"``]
the metric used to compute surface distance. Defaults to ``"euclidean"``.
"""
if not include_background:
y_pred, y = ignore_background(y_pred=y_pred, y=y)
if isinstance(y, torch.Tensor):
y = y.float()
if isinstance(y_pred, torch.Tensor):
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError(
f"y_pred and y should have same shapes, got {y_pred.shape} and {y.shape}."
)
batch_size, n_class = y_pred.shape[:2]
asd = np.empty((batch_size, n_class))
for b, c in np.ndindex(batch_size, n_class):
(edges_pred, edges_gt) = get_mask_edges(y_pred[b, c], y[b, c])
if not np.any(edges_gt):
warnings.warn(
f"the ground truth of class {c} is all 0, this may result in nan/inf distance."
)
if not np.any(edges_pred):
warnings.warn(
f"the prediction of class {c} is all 0, this may result in nan/inf distance."
)
surface_distance = get_surface_distance(
edges_pred, edges_gt, distance_metric=distance_metric)
if symmetric:
surface_distance_2 = get_surface_distance(
edges_gt, edges_pred, distance_metric=distance_metric)
surface_distance = np.concatenate(
[surface_distance, surface_distance_2])
asd[b, c] = np.nan if surface_distance.shape == (
0, ) else surface_distance.mean()
return convert_data_type(asd, torch.Tensor)[0]
def compute_average_surface_distance(
y_pred: Union[np.ndarray, torch.Tensor],
y: Union[np.ndarray, torch.Tensor],
include_background: bool = False,
symmetric: bool = False,
distance_metric: str = "euclidean",
):
"""
This function is used to compute the Average Surface Distance from `y_pred` to `y`
under the default setting.
In addition, if sets ``symmetric = True``, the average symmetric surface distance between
these two inputs will be returned.
Args:
y_pred: input data to compute, typical segmentation model output.
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
should be binarized.
y: ground truth to compute mean the distance. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip distance computation on the first channel of
the predicted output. Defaults to ``False``.
symmetric: whether to calculate the symmetric average surface distance between
`seg_pred` and `seg_gt`. Defaults to ``False``.
distance_metric: : [``"euclidean"``, ``"chessboard"``, ``"taxicab"``]
the metric used to compute surface distance. Defaults to ``"euclidean"``.
"""
if not include_background:
y_pred, y = ignore_background(
y_pred=y_pred,
y=y,
)
y = y.float()
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError("y_pred and y should have same shapes.")
batch_size, n_class = y_pred.shape[:2]
asd = np.empty((batch_size, n_class))
for b, c in np.ndindex(batch_size, n_class):
(edges_pred, edges_gt) = get_mask_edges(y_pred[b, c], y[b, c])
surface_distance = get_surface_distance(
edges_pred, edges_gt, distance_metric=distance_metric)
if surface_distance.shape == (0, ):
avg_surface_distance = np.nan
else:
avg_surface_distance = surface_distance.mean()
if not symmetric:
asd[b, c] = avg_surface_distance
else:
surface_distance_2 = get_surface_distance(
edges_gt, edges_pred, distance_metric=distance_metric)
if surface_distance_2.shape == (0, ):
avg_surface_distance_2 = np.nan
else:
avg_surface_distance_2 = surface_distance_2.mean()
asd[b, c] = np.mean((avg_surface_distance, avg_surface_distance_2))
return torch.from_numpy(asd)
def compute_generalized_dice(
y_pred: torch.Tensor,
y: torch.Tensor,
include_background: bool = True,
weight_type: Union[Weight, str] = Weight.SQUARE,
) -> torch.Tensor:
"""Computes the Generalized Dice Score and returns a tensor with its per image values.
Args:
y_pred (torch.Tensor): binarized segmentation model output. It should be binarized, in one-hot format
and in the NCHW[D] format, where N is the batch dimension, C is the channel dimension, and the
remaining are the spatial dimensions.
y (torch.Tensor): binarized ground-truth. It should be binarized, in one-hot format and have the same shape as `y_pred`.
include_background (bool, optional): whether to skip score computation on the first channel of the
predicted output. Defaults to True.
weight_type (Union[Weight, str], optional): {``"square"``, ``"simple"``, ``"uniform"``}. Type of function to
transform ground truth volume into a weight factor. Defaults to ``"square"``.
Returns:
torch.Tensor: per batch and per class Generalized Dice Score, i.e., with the shape [batch_size, num_classes].
Raises:
ValueError: if `y_pred` or `y` are not PyTorch tensors, if `y_pred` and `y` have less than three dimensions,
or `y_pred` and `y` don't have the same shape.
"""
# Ensure tensors are binarized
is_binary_tensor(y_pred, "y_pred")
is_binary_tensor(y, "y")
# Ensure tensors have at least 3 dimensions and have the same shape
dims = y_pred.dim()
if dims < 3:
raise ValueError(
f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}."
)
if y.shape != y_pred.shape:
raise ValueError(
f"y_pred - {y_pred.shape} - and y - {y.shape} - should have the same shapes."
)
# Ignore background, if needed
if not include_background:
y_pred, y = ignore_background(y_pred=y_pred, y=y)
# Reducing only spatial dimensions (not batch nor channels), compute the intersection and non-weighted denominator
reduce_axis = list(range(2, y_pred.dim()))
intersection = torch.sum(y * y_pred, dim=reduce_axis)
y_o = torch.sum(y, dim=reduce_axis)
y_pred_o = torch.sum(y_pred, dim=reduce_axis)
denominator = y_o + y_pred_o
# Set the class weights
weight_type = look_up_option(weight_type, Weight)
if weight_type == Weight.SIMPLE:
w = torch.reciprocal(y_o.float())
elif weight_type == Weight.SQUARE:
w = torch.reciprocal(y_o.float() * y_o.float())
else:
w = torch.ones_like(y_o.float())
# Replace infinite values for non-appearing classes by the maximum weight
for b in w:
infs = torch.isinf(b)
b[infs] = 0
b[infs] = torch.max(b)
# Compute the weighted numerator and denominator, summing along the class axis
numer = 2.0 * (intersection * w).sum(dim=1)
denom = (denominator * w).sum(dim=1)
# Compute the score
generalized_dice_score = numer / denom
# Handle zero deivision. Where denom == 0 and the prediction volume is 0, score is 1.
# Where denom == 0 but the prediction volume is not 0, score is 0
y_pred_o = y_pred_o.sum(dim=-1)
denom_zeros = denom == 0
generalized_dice_score[denom_zeros] = torch.where(
(y_pred_o == 0)[denom_zeros], torch.tensor(1.0), torch.tensor(0.0))
return generalized_dice_score
def compute_surface_dice(
y_pred: torch.Tensor,
y: torch.Tensor,
class_thresholds: List[float],
include_background: bool = False,
distance_metric: str = "euclidean",
):
r"""
This function computes the (Normalized) Surface Dice (NSD) between the two tensors `y_pred` (referred to as
:math:`\hat{Y}`) and `y` (referred to as :math:`Y`). This metric determines which fraction of a segmentation
boundary is correctly predicted. A boundary element is considered correctly predicted if the closest distance to the
reference boundary is smaller than or equal to the specified threshold related to the acceptable amount of deviation in
pixels. The NSD is bounded between 0 and 1.
This implementation supports multi-class tasks with an individual threshold :math:`\tau_c` for each class :math:`c`.
The class-specific NSD for batch index :math:`b`, :math:`\operatorname {NSD}_{b,c}`, is computed using the function:
.. math::
\operatorname {NSD}_{b,c} \left(Y_{b,c}, \hat{Y}_{b,c}\right) = \frac{\left|\mathcal{D}_{Y_{b,c}}^{'}\right| +
\left| \mathcal{D}_{\hat{Y}_{b,c}}^{'} \right|}{\left|\mathcal{D}_{Y_{b,c}}\right| +
\left|\mathcal{D}_{\hat{Y}_{b,c}}\right|}
:label: nsd
with :math:`\mathcal{D}_{Y_{b,c}}` and :math:`\mathcal{D}_{\hat{Y}_{b,c}}` being two sets of nearest-neighbor
distances. :math:`\mathcal{D}_{Y_{b,c}}` is computed from the predicted segmentation boundary towards the reference segmentation
boundary and vice-versa for :math:`\mathcal{D}_{\hat{Y}_{b,c}}`. :math:`\mathcal{D}_{Y_{b,c}}^{'}` and
:math:`\mathcal{D}_{\hat{Y}_{b,c}}^{'}` refer to the subsets of distances that are smaller or equal to the
acceptable distance :math:`\tau_c`:
.. math::
\mathcal{D}_{Y_{b,c}}^{'} = \{ d \in \mathcal{D}_{Y_{b,c}} \, | \, d \leq \tau_c \}.
In the case of a class neither being present in the predicted segmentation, nor in the reference segmentation, a nan value
will be returned for this class. In the case of a class being present in only one of predicted segmentation or
reference segmentation, the class NSD will be 0.
This implementation is based on https://arxiv.org/abs/2111.05408 and supports 2D images.
Be aware that the computation of boundaries is different from DeepMind's implementation
https://github.com/deepmind/surface-distance. In this implementation, the length of a segmentation boundary is
interpreted as the number of its edge pixels. In DeepMind's implementation, the length of a segmentation boundary
depends on the local neighborhood (cf. https://arxiv.org/abs/1809.04430).
Args:
y_pred: Predicted segmentation, typically segmentation model output.
It must be a one-hot encoded, batch-first tensor [B,C,H,W].
y: Reference segmentation.
It must be a one-hot encoded, batch-first tensor [B,C,H,W].
class_thresholds: List of class-specific thresholds.
The thresholds relate to the acceptable amount of deviation in the segmentation boundary in pixels.
Each threshold needs to be a finite, non-negative number.
include_background: Whether to skip the surface dice computation on the first channel of
the predicted output. Defaults to ``False``.
distance_metric: The metric used to compute surface distances.
One of [``"euclidean"``, ``"chessboard"``, ``"taxicab"``].
Defaults to ``"euclidean"``.
Raises:
ValueError: If `y_pred` and/or `y` are not PyTorch tensors.
ValueError: If `y_pred` and/or `y` do not have four dimensions.
ValueError: If `y_pred` and/or `y` have different shapes.
ValueError: If `y_pred` and/or `y` are not one-hot encoded
ValueError: If the number of channels of `y_pred` and/or `y` is different from the number of class thresholds.
ValueError: If any class threshold is not finite.
ValueError: If any class threshold is negative.
Returns:
Pytorch Tensor of shape [B,C], containing the NSD values :math:`\operatorname {NSD}_{b,c}` for each batch index
:math:`b` and class :math:`c`.
"""
if not include_background:
y_pred, y = ignore_background(y_pred=y_pred, y=y)
if not isinstance(y_pred, torch.Tensor) or not isinstance(y, torch.Tensor):
raise ValueError("y_pred and y must be PyTorch Tensor.")
if y_pred.ndimension() != 4 or y.ndimension() != 4:
raise ValueError(
"y_pred and y should have four dimensions: [B,C,H,W].")
if y_pred.shape != y.shape:
raise ValueError(
f"y_pred and y should have same shape, but instead, shapes are {y_pred.shape} (y_pred) and {y.shape} (y)."
)
if not torch.all(y_pred.byte() == y_pred) or not torch.all(y.byte() == y):
raise ValueError(
"y_pred and y should be binarized tensors (e.g. torch.int64).")
if torch.any(y_pred > 1) or torch.any(y > 1):
raise ValueError("y_pred and y should be one-hot encoded.")
y = y.float()
y_pred = y_pred.float()
batch_size, n_class = y_pred.shape[:2]
if n_class != len(class_thresholds):
raise ValueError(
f"number of classes ({n_class}) does not match number of class thresholds ({len(class_thresholds)})."
)
if any(~np.isfinite(class_thresholds)):
raise ValueError("All class thresholds need to be finite.")
if any(np.array(class_thresholds) < 0):
raise ValueError("All class thresholds need to be >= 0.")
nsd = np.empty((batch_size, n_class))
for b, c in np.ndindex(batch_size, n_class):
(edges_pred, edges_gt) = get_mask_edges(y_pred[b, c],
y[b, c],
crop=False)
if not np.any(edges_gt):
warnings.warn(
f"the ground truth of class {c} is all 0, this may result in nan/inf distance."
)
if not np.any(edges_pred):
warnings.warn(
f"the prediction of class {c} is all 0, this may result in nan/inf distance."
)
distances_pred_gt = get_surface_distance(
edges_pred, edges_gt, distance_metric=distance_metric)
distances_gt_pred = get_surface_distance(
edges_gt, edges_pred, distance_metric=distance_metric)
boundary_complete = len(distances_pred_gt) + len(distances_gt_pred)
boundary_correct = np.sum(
distances_pred_gt <= class_thresholds[c]) + np.sum(
distances_gt_pred <= class_thresholds[c])
if boundary_complete == 0:
# the class is neither present in the prediction, nor in the reference segmentation
nsd[b, c] = np.nan
else:
nsd[b, c] = boundary_correct / boundary_complete
return convert_data_type(nsd, torch.Tensor)[0]
