| from __future__ import annotations |
| from typing import Dict, List |
| import numpy as np |
| import torch |
| from scipy import ndimage |
|
|
|
|
| def dice_per_class(pred: np.ndarray, target: np.ndarray, num_classes: int, include_bg: bool = False) -> Dict[str, float]: |
| out = {} |
| start = 0 if include_bg else 1 |
| vals = [] |
| for c in range(start, num_classes): |
| p = pred == c |
| t = target == c |
| denom = p.sum() + t.sum() |
| if denom == 0: |
| d = np.nan |
| else: |
| d = 2.0 * np.logical_and(p, t).sum() / denom |
| out[f"dice_c{c}"] = float(d) if not np.isnan(d) else np.nan |
| if not np.isnan(d): |
| vals.append(d) |
| out["dice_mean"] = float(np.mean(vals)) if vals else np.nan |
| return out |
|
|
|
|
| def _surface(mask: np.ndarray) -> np.ndarray: |
| if mask.sum() == 0: |
| return mask.astype(bool) |
| eroded = ndimage.binary_erosion(mask, iterations=1, border_value=0) |
| return np.logical_xor(mask, eroded) |
|
|
|
|
| def hd95_per_class(pred: np.ndarray, target: np.ndarray, num_classes: int, spacing=None, include_bg: bool = False, empty_penalty: str = "diagonal") -> Dict[str, float]: |
| """Compute per-class HD95. |
| |
| If one mask is empty and the other is not, assigning NaN and dropping the |
| class makes failed segmentations look artificially good. We instead use the |
| physical image diagonal as a conservative finite penalty. If both masks are |
| empty, the class is not present and is excluded from the mean. |
| """ |
| out = {} |
| start = 0 if include_bg else 1 |
| vals = [] |
| if spacing is None: |
| spacing = (1.0, 1.0, 1.0) |
| spacing = tuple(float(x) for x in spacing) |
| diag = float(np.sqrt(sum(((s * max(1, n - 1)) ** 2) for s, n in zip(spacing, target.shape)))) |
| for c in range(start, num_classes): |
| p = pred == c |
| t = target == c |
| p_sum = int(p.sum()) |
| t_sum = int(t.sum()) |
| if p_sum == 0 and t_sum == 0: |
| hd = np.nan |
| elif p_sum == 0 or t_sum == 0: |
| hd = diag if empty_penalty == "diagonal" else np.nan |
| else: |
| ps = _surface(p) |
| ts = _surface(t) |
| if ps.sum() == 0 or ts.sum() == 0: |
| hd = diag if empty_penalty == "diagonal" else np.nan |
| else: |
| dt_t = ndimage.distance_transform_edt(~ts, sampling=spacing) |
| dt_p = ndimage.distance_transform_edt(~ps, sampling=spacing) |
| dists = np.concatenate([dt_t[ps], dt_p[ts]]) |
| hd = np.percentile(dists, 95) if dists.size else np.nan |
| out[f"hd95_c{c}"] = float(hd) if not np.isnan(hd) else np.nan |
| if not np.isnan(hd): |
| vals.append(hd) |
| out["hd95_mean"] = float(np.mean(vals)) if vals else np.nan |
| return out |
|
|
|
|
| def torch_soft_dice_loss(logits: torch.Tensor, target: torch.Tensor, num_classes: int, ignore_index: int | None = None, eps: float = 1e-5): |
| probs = torch.softmax(logits, dim=1) |
| if target.ndim == logits.ndim: |
| onehot = target.float() |
| else: |
| target_clamped = target.clamp(0, num_classes - 1).long() |
| onehot = torch.nn.functional.one_hot(target_clamped, num_classes).permute(0, 4, 1, 2, 3).float() |
| if ignore_index is not None: |
| mask = target != ignore_index |
| probs = probs * mask.unsqueeze(1) |
| onehot = onehot * mask.unsqueeze(1) |
| dims = tuple(range(2, logits.ndim)) |
| inter = (probs * onehot).sum(dims) |
| denom = probs.sum(dims) + onehot.sum(dims) |
| dice = (2 * inter + eps) / (denom + eps) |
| return 1.0 - dice[:, 1:].mean() |
|
|
|
|
| def entropy_loss(logits: torch.Tensor, eps: float = 1e-8): |
| p = torch.softmax(logits, dim=1) |
| return -(p * (p + eps).log()).sum(dim=1).mean() |
|
|
|
|
| def confidence_and_margin(probs: torch.Tensor): |
| vals, inds = probs.topk(k=2, dim=1) |
| conf = vals[:, 0] |
| margin = vals[:, 0] - vals[:, 1] |
| pred = inds[:, 0] |
| return conf, margin, pred |
|
|
|
|
| def finite_difference_boundary(x: torch.Tensor) -> torch.Tensor: |
| |
| dx = torch.abs(x[:, :, 1:] - x[:, :, :-1]).mean(dim=1, keepdim=True) |
| dx = torch.nn.functional.pad(dx, (0,0,0,0,1,0)) |
| dy = torch.abs(x[:, :, :, 1:] - x[:, :, :, :-1]).mean(dim=1, keepdim=True) |
| dy = torch.nn.functional.pad(dy, (0,0,1,0,0,0)) |
| dz = torch.abs(x[:, :, :, :, 1:] - x[:, :, :, :, :-1]).mean(dim=1, keepdim=True) |
| dz = torch.nn.functional.pad(dz, (1,0,0,0,0,0)) |
| return (dx + dy + dz) / 3.0 |
|
|
