Restore statistical padding modes

src/torchio/data/sampler.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterator
+from typing import TYPE_CHECKING
 from typing import Any
 
 import numpy as np
@@ -15,6 +16,9 @@
 from .patch import PatchLocation
 from .subject import Subject
 
+if TYPE_CHECKING:
+    from ..transforms.spatial._padding import PaddingMode
+
 
 class PatchSampler:
     """Base class for patch samplers.
@@ -95,7 +99,7 @@ def __init__(
         subject: Subject,
         patch_size: int | TypeThreeInts,
         patch_overlap: int | TypeThreeInts = 0,
-        padding_mode: str | None = None,
+        padding_mode: PaddingMode | None = None,
         fill: float = 0,
     ) -> None:
         super().__init__(patch_size)

src/torchio/transforms/_statistics.py

@@ -0,0 +1,43 @@
+"""Shared statistical helpers for transforms."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor
+
+
+def compute_quantile(values: Tensor, q: float) -> Tensor:
+    """Compute a single quantile of a one-dimensional tensor.
+
+    `torch.quantile` raises `RuntimeError: quantile() input tensor is too
+    large` for inputs with more than `2**24` elements, which is easily
+    reached by high-resolution volumes. `torch.kthvalue` has no such size
+    limit and is much faster on large tensors, so it is used here with linear
+    interpolation to reproduce the default `torch.quantile` behavior.
+
+    This is adapted from the solution by Elie Goudout (`@ego-thales`):
+    https://github.com/pytorch/pytorch/issues/157431#issuecomment-3026856373
+
+    Args:
+        values: One-dimensional tensor of values.
+        q: Quantile to compute, in the `[0, 1]` range.
+
+    Returns:
+        Zero-dimensional tensor with the computed quantile.
+
+    Raises:
+        ValueError: If `q` is outside the `[0, 1]` range.
+    """
+    if not 0 <= q <= 1:
+        msg = f"Only values 0 <= q <= 1 are supported, but got {q!r}"
+        raise ValueError(msg)
+    index = q * (values.numel() - 1)
+    lower = math.floor(index)
+    lower_value = torch.kthvalue(values, lower + 1).values
+    if index == lower:
+        return lower_value
+    upper_value = torch.kthvalue(values, lower + 2).values
+    weight = index - lower
+    return lower_value.lerp(upper_value, weight)

src/torchio/transforms/intensity/normalize.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import math
 import warnings
 from collections.abc import Callable
 from typing import Any
@@ -14,6 +13,7 @@
 from ...data.batch import ImagesBatch
 from ...data.batch import SubjectsBatch
 from ...data.image import LabelMap
+from .._statistics import compute_quantile
 from ..parameter_range import Choice
 from ..parameter_range import _ParameterRange
 from ..transform import IntensityTransform
@@ -360,45 +360,10 @@ def _percentile_range(
         )
         values = tensor.reshape(-1)
 
-    low = float(_quantile(values.float(), pct_low / 100.0).item())
-    high = float(_quantile(values.float(), pct_high / 100.0).item())
+    low = float(compute_quantile(values.float(), pct_low / 100.0).item())
+    high = float(compute_quantile(values.float(), pct_high / 100.0).item())
     return low, high
 
 
-def _quantile(values: Tensor, q: float) -> Tensor:
-    """Compute a single quantile of a 1D tensor using ``torch.kthvalue``.
-
-    ``torch.quantile`` raises ``RuntimeError: quantile() input tensor is too
-    large`` for inputs with more than ``2**24`` elements, which is easily
-    reached by high-resolution volumes. ``torch.kthvalue`` has no such size
-    limit and is much faster on large tensors, so it is used here with linear
-    interpolation to reproduce the default ``torch.quantile`` behaviour.
-
-    This is adapted from the solution by Élie Goudout (``@ego-thales``):
-    https://github.com/pytorch/pytorch/issues/157431#issuecomment-3026856373
-
-    Args:
-        values: One-dimensional tensor of values.
-        q: Quantile to compute, in the ``[0, 1]`` range.
-
-    Returns:
-        Zero-dimensional tensor with the computed quantile.
-
-    Raises:
-        ValueError: If ``q`` is outside the ``[0, 1]`` range.
-    """
-    if not 0 <= q <= 1:
-        msg = f"Only values 0 <= q <= 1 are supported, but got {q!r}"
-        raise ValueError(msg)
-    index = q * (values.numel() - 1)
-    lower = math.floor(index)
-    lower_value = torch.kthvalue(values, lower + 1).values
-    if index == lower:
-        return lower_value
-    upper_value = torch.kthvalue(values, lower + 2).values
-    weight = index - lower
-    return lower_value.lerp(upper_value, weight)
-
-
 # Backwards-compatible alias.
 RescaleIntensity = Normalize

src/torchio/transforms/spatial/_padding.py

@@ -0,0 +1,99 @@
+"""Shared spatial padding helpers."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Literal
+from typing import get_args
+
+import torch
+from torch import Tensor
+
+from ...types import TypeSixInts
+from .._statistics import compute_quantile
+
+#: Accepted padding modes.
+PaddingMode = Literal[
+    "constant",
+    "reflect",
+    "replicate",
+    "circular",
+    "mean",
+    "median",
+    "minimum",
+]
+
+_PADDING_MODES: tuple[PaddingMode, ...] = get_args(PaddingMode)
+_STATISTIC_PADDING_MODES = "mean", "median", "minimum"
+
+
+def parse_padding_mode(padding_mode: str) -> PaddingMode:
+    """Validate and return a padding mode."""
+    if padding_mode not in _PADDING_MODES:
+        msg = f"padding_mode must be one of {_PADDING_MODES}, got {padding_mode!r}"
+        raise ValueError(msg)
+    return padding_mode
+
+
+def _compute_padding_statistic(
+    data: Tensor,
+    padding_mode: PaddingMode,
+) -> Tensor:
+    """Compute one whole-volume padding statistic per batch element."""
+    flat = data.flatten(start_dim=1)
+    if padding_mode == "minimum":
+        return flat.amin(dim=1)
+
+    if not torch.is_floating_point(data):
+        warnings.warn(
+            f'The constant value computed for padding mode "{padding_mode}"'
+            " might be truncated in the output, as the data type of the input"
+            " image is not float. Consider converting the image to a floating"
+            " point type before applying this transform.",
+            RuntimeWarning,
+            stacklevel=4,
+        )
+
+    float_flat = flat if data.dtype in (torch.float32, torch.float64) else flat.float()
+    if padding_mode == "mean":
+        statistic = float_flat.mean(dim=1)
+    else:
+        statistic = torch.stack(
+            [compute_quantile(values, 0.5) for values in float_flat],
+        )
+    return statistic.to(data.dtype)
+
+
+def pad_tensor(
+    data: Tensor,
+    padding: TypeSixInts,
+    padding_mode: PaddingMode,
+    fill: float,
+) -> Tensor:
+    """Pad a 4D image tensor or 5D image batch."""
+    if data.ndim not in (4, 5):
+        msg = f"Expected a 4D or 5D image tensor, got {data.ndim}D"
+        raise ValueError(msg)
+    i0, i1, j0, j1, k0, k1 = padding
+    pad_arg = k0, k1, j0, j1, i0, i1
+    if padding_mode not in _STATISTIC_PADDING_MODES:
+        return torch.nn.functional.pad(
+            data,
+            pad_arg,
+            mode=padding_mode,
+            value=fill,
+        )
+
+    is_unbatched = data.ndim == 4
+    batch = data.unsqueeze(0) if is_unbatched else data
+    statistic = _compute_padding_statistic(batch, padding_mode)
+    padded = torch.nn.functional.pad(batch, pad_arg)
+    interior = torch.ones(
+        (1, 1, *batch.shape[-3:]),
+        dtype=torch.bool,
+        device=batch.device,
+    )
+    interior = torch.nn.functional.pad(interior, pad_arg)
+    fill_values = statistic.reshape(-1, 1, 1, 1, 1)
+    result = torch.where(interior, padded, fill_values)
+    return result[0] if is_unbatched else result

src/torchio/transforms/spatial/crop_or_pad.py

@@ -27,6 +27,9 @@
 from ..compose import Compose
 from ..transform import AppliedTransform
 from ..transform import SpatialTransform
+from ._padding import PaddingMode
+from ._padding import pad_tensor
+from ._padding import parse_padding_mode
 from .crop import Crop
 from .pad import Pad
 
@@ -208,7 +211,7 @@ def __init__(
         padding: TypeSixInts,
         padded_shape: tuple[int, int, int, int],
         affine: TypeAffineMatrix,
-        padding_mode: str = "constant",
+        padding_mode: PaddingMode = "constant",
         fill: float = 0,
     ) -> None:
         self._source = source
@@ -234,13 +237,11 @@ def dtype(self) -> np.dtype:
 
     def to_tensor(self) -> Tensor:
         base = self._source.to_tensor()
-        i0, i1, j0, j1, k0, k1 = self._padding
-        pad_arg = (k0, k1, j0, j1, i0, i1)
-        return torch.nn.functional.pad(
+        return pad_tensor(
             base,
-            pad_arg,
-            mode=self._padding_mode,
-            value=self._fill,
+            self._padding,
+            self._padding_mode,
+            self._fill,
         )
 
     def __getitem__(self, slices: SliceIndex) -> Tensor:
@@ -318,7 +319,7 @@ def _crop_image_lazy(image: Image, cropping: TypeSixInts) -> Image:
 def _pad_image_lazy(
     image: Image,
     padding: TypeSixInts,
-    padding_mode: str,
+    padding_mode: PaddingMode,
     fill: float,
 ) -> Image:
     """Pad an image lazily (data is only loaded when accessed)."""
@@ -336,12 +337,11 @@ def _pad_image_lazy(
     padded_shape = (c, si + i0 + i1, sj + j0 + j1, sk + k0 + k1)
 
     if image.is_loaded:
-        pad_arg = (k0, k1, j0, j1, i0, i1)
-        new_data = torch.nn.functional.pad(
+        new_data = pad_tensor(
             image.data,
-            pad_arg,
-            mode=padding_mode,
-            value=fill,
+            padding,
+            padding_mode,
+            fill,
         )
         return image.new_like(data=new_data, affine=new_affine)
 
@@ -369,12 +369,11 @@ def _pad_image_lazy(
         return new
 
     # No backend → fall back to eager pad
-    pad_arg = (k0, k1, j0, j1, i0, i1)
-    new_data = torch.nn.functional.pad(
+    new_data = pad_tensor(
         image.data,
-        pad_arg,
-        mode=padding_mode,
-        value=fill,
+        padding,
+        padding_mode,
+        fill,
     )
     return image.new_like(data=new_data, affine=new_affine)
 
@@ -403,8 +402,11 @@ class CropOrPad(SpatialTransform):
         units: Coordinate system for `target_shape`. One of
             `"voxels"` (default), `"mm"`, or `"cm"`.
         padding_mode: One of `'constant'`, `'reflect'`,
-            `'replicate'`, or `'circular'`. See
-            [`torch.nn.functional.pad`](https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html).
+            `'replicate'`, `'circular'`, `'mean'`, `'median'`, or
+            `'minimum'`. Statistical modes use one value computed from
+            the whole image volume. For integer inputs, `'mean'` and
+            `'median'` may be truncated to the input dtype and emit a
+            warning.
         fill: Fill value when `padding_mode='constant'`.
         only_crop: If `True`, padding is never applied. Mutually
             exclusive with `only_pad`.
@@ -426,14 +428,15 @@ class CropOrPad(SpatialTransform):
         >>> transform = tio.CropOrPad(target_shape=256, only_pad=True)
         >>> transform = tio.CropOrPad(target_shape=(256, 256, None))  # keep depth
         >>> transform = tio.CropOrPad(target_shape=96, location='random')
+        >>> transform = tio.CropOrPad(target_shape=256, padding_mode='mean')
     """
 
     def __init__(
         self,
         target_shape: TargetShapeParam,
         *,
         units: Units = "voxels",
-        padding_mode: str = "constant",
+        padding_mode: PaddingMode = "constant",
         fill: float = 0,
         only_crop: bool = False,
         only_pad: bool = False,
@@ -452,7 +455,7 @@ def __init__(
             raise ValueError(msg)
         self.target_shape = _parse_target_shape(target_shape)
         self.units: Units = units
-        self.padding_mode = padding_mode
+        self.padding_mode = parse_padding_mode(padding_mode)
         self.fill = fill
         self.only_crop = only_crop
         self.only_pad = only_pad

src/torchio/transforms/spatial/ensure_shape_multiple.py

@@ -10,6 +10,8 @@
 from ...data.subject import Subject
 from ...types import TypeThreeInts
 from ..transform import SpatialTransform
+from ._padding import PaddingMode
+from ._padding import parse_padding_mode
 from .crop_or_pad import CropOrPad
 
 #: Accepted target_multiple specifications.
@@ -74,7 +76,8 @@ class EnsureShapeMultiple(SpatialTransform):
             multiple.
         padding_mode: Padding mode forwarded to `CropOrPad` when
             `method='pad'`. One of `'constant'`, `'reflect'`,
-            `'replicate'`, or `'circular'`.
+            `'replicate'`, `'circular'`, `'mean'`, `'median'`, or
+            `'minimum'`.
         fill: Fill value when `padding_mode='constant'`.
         **kwargs: See [`Transform`][torchio.Transform] for additional
             keyword arguments.
@@ -92,7 +95,7 @@ def __init__(
         target_multiple: TargetMultipleParam,
         *,
         method: str = "pad",
-        padding_mode: str = "constant",
+        padding_mode: PaddingMode = "constant",
         fill: float = 0,
         **kwargs: Any,
     ) -> None:
@@ -102,7 +105,7 @@ def __init__(
             msg = f"method must be 'crop' or 'pad', got {method!r}"
             raise ValueError(msg)
         self.method = method
-        self.padding_mode = padding_mode
+        self.padding_mode = parse_padding_mode(padding_mode)
         self.fill = fill
 
     def forward(self, data: Any) -> Any: