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/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 ..intensity.normalize import _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(
+            [_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,7 @@
 from ...data.subject import Subject
 from ...types import TypeThreeInts
 from ..transform import SpatialTransform
+from ._padding import PaddingMode
 from .crop_or_pad import CropOrPad
 
 #: Accepted target_multiple specifications.
@@ -74,7 +75,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 +94,7 @@ def __init__(
         target_multiple: TargetMultipleParam,
         *,
         method: str = "pad",
-        padding_mode: str = "constant",
+        padding_mode: PaddingMode = "constant",
         fill: float = 0,
         **kwargs: Any,
     ) -> None:

src/torchio/transforms/spatial/pad.py

@@ -4,12 +4,13 @@
 
 from typing import Any
 
-import torch
-
 from ...data.batch import SubjectsBatch
 from ...types import TypeSixInts
 from ...types import TypeThreeInts
 from ..transform import SpatialTransform
+from ._padding import PaddingMode
+from ._padding import pad_tensor
+from ._padding import parse_padding_mode
 
 #: Accepted padding specifications.
 #: `int` → same amount on each side of each axis.
@@ -50,8 +51,11 @@ class Pad(SpatialTransform):
             $i_\text{ini} = i_\text{fin} = i$, etc.
             If only one value $n$ is provided, all six values are $n$.
         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'`.
         **kwargs: See [`Transform`][torchio.Transform] for additional
             keyword arguments.
@@ -61,19 +65,20 @@ class Pad(SpatialTransform):
         >>> transform = tio.Pad(padding=10)
         >>> transform = tio.Pad(padding=(5, 10, 0))
         >>> transform = tio.Pad(padding=10, padding_mode='reflect')
+        >>> transform = tio.Pad(padding=10, padding_mode='minimum')
     """
 
     def __init__(
         self,
         *,
         padding: PaddingParam,
-        padding_mode: str = "constant",
+        padding_mode: PaddingMode = "constant",
         fill: float = 0,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.padding = _parse_padding(padding)
-        self.padding_mode = padding_mode
+        self.padding_mode = parse_padding_mode(padding_mode)
         self.fill = fill
 
     def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
@@ -91,14 +96,12 @@ def apply_transform(
         i0, i1, j0, j1, k0, k1 = params["padding"]
         mode = params["padding_mode"]
         fill = params["fill"]
-        # F.pad expects reversed order: (k0, k1, j0, j1, i0, i1)
-        pad_arg = (k0, k1, j0, j1, i0, i1)
         for _name, img_batch in self._get_images(batch).items():
-            img_batch.data = torch.nn.functional.pad(
+            img_batch.data = pad_tensor(
                 img_batch.data,
-                pad_arg,
-                mode=mode,
-                value=fill,
+                (i0, i1, j0, j1, k0, k1),
+                mode,
+                fill,
             )
             # Update each affine's origin (shift back)
             for affine in img_batch.affines: