Use global statistics in Pad (and CropOrPad) (#1336)

Author: fepegar

src/torchio/transforms/preprocessing/spatial/pad.py

@@ -2,11 +2,11 @@
 from numbers import Number
 from typing import Union
 
-import nibabel as nib
 import numpy as np
 import torch
+from nibabel.affines import apply_affine
 
-from ....data.image import LabelMap
+from ....data.image import Image
 from ....data.subject import Subject
 from .bounds_transform import BoundsTransform
 from .bounds_transform import TypeBounds
@@ -31,7 +31,10 @@ class Pad(BoundsTransform):
             :math:`w_{ini} = w_{fin} = h_{ini} = h_{fin} =
             d_{ini} = d_{fin} = n`.
         padding_mode: See possible modes in `NumPy docs`_. If it is a number,
-            the mode will be set to ``'constant'``.
+            the mode will be set to ``'constant'``. If it is ``'mean'``,
+            ``'maximum'``, ``'median'`` or ``'minimum'``, the statistic will be
+            computed from the whole volume, unlike in NumPy, which computes it
+            along the padded axis.
         **kwargs: See :class:`~torchio.transforms.Transform` for additional
             keyword arguments.
 
@@ -78,26 +81,49 @@ def check_padding_mode(cls, padding_mode):
             )
             raise KeyError(message)
 
+    def _check_truncation(self, image: Image, mode: Union[str, float]) -> None:
+        if mode not in ('mean', 'median'):
+            return
+        if torch.is_floating_point(image.data):
+            return
+        message = (
+            f'The constant value computed for padding mode "{mode}" might '
+            ' be truncated in the output, as the input image is not'
+            'floating point. Consider converting the image to a floating'
+            ' point type before applying this transform.'
+        )
+        warnings.warn(message, RuntimeWarning, stacklevel=2)
+
     def apply_transform(self, subject: Subject) -> Subject:
         assert self.bounds_parameters is not None
         low = self.bounds_parameters[::2]
         for image in self.get_images(subject):
-            if isinstance(image, LabelMap) and self.padding_mode == 'mean':
-                message = (
-                    'Padding mode "mean" might create non-integer values in label maps'
-                )
-                warnings.warn(message, RuntimeWarning, stacklevel=2)
-            new_origin = nib.affines.apply_affine(image.affine, -np.array(low))
+            self._check_truncation(image, self.padding_mode)
+            new_origin = apply_affine(image.affine, -np.array(low))
             new_affine = image.affine.copy()
             new_affine[:3, 3] = new_origin
-            kwargs: dict[str, Union[str, float]]
+
+            mode: str | float = 'constant'
+            constant: torch.Tensor | float | None = None
+            kwargs: dict[str, str | float | torch.Tensor] = {}
             if isinstance(self.padding_mode, Number):
-                kwargs = {
-                    'mode': 'constant',
-                    'constant_values': self.padding_mode,
-                }
+                constant = self.padding_mode  # type: ignore[assignment]
+            elif self.padding_mode == 'maximum':
+                constant = image.data.max()
+            elif self.padding_mode == 'mean':
+                constant = image.data.float().mean()
+            elif self.padding_mode == 'median':
+                constant = torch.quantile(image.data.float(), 0.5)
+            elif self.padding_mode == 'minimum':
+                constant = image.data.min()
             else:
-                kwargs = {'mode': self.padding_mode}
+                constant = None
+                mode = self.padding_mode
+
+            if constant is not None:
+                kwargs['constant_values'] = constant
+            kwargs['mode'] = mode
+
             pad_params = self.bounds_parameters
             paddings = (0, 0), pad_params[:2], pad_params[2:4], pad_params[4:]
             padded = np.pad(image.data, paddings, **kwargs)  # type: ignore[call-overload]

tests/transforms/preprocessing/test_pad.py

@@ -1,3 +1,4 @@
+import pytest
 import SimpleITK as sitk
 import torch
 
@@ -40,3 +41,38 @@ def padding_func():
     def test_padding_mean_label_map(self):
         with self.assertWarns(RuntimeWarning):
             tio.Pad(1, padding_mode='mean')(self.sample_subject.label)
+
+    def test_padding_modes_global(self):
+        x = torch.ones(1, 1, 2, 2, dtype=torch.int)
+        x[..., 0, 0] = 0
+        # The image should look like this:
+        # 0 1
+        # 1 1
+
+        add_bottom_row = 0, 0, 0, 1, 0, 0
+        with_zeros = tio.Pad(add_bottom_row)(x)
+        assert with_zeros[0, 0, 2].tolist() == [0, 0]
+
+        with_minimum = tio.Pad(add_bottom_row, padding_mode='minimum')(x)
+        assert with_minimum[0, 0, 2].tolist() == [0, 0]
+
+        with_maximum = tio.Pad(add_bottom_row, padding_mode='maximum')(x)
+        assert with_maximum[0, 0, 2].tolist() == [1, 1]
+
+        with_median = tio.Pad(add_bottom_row, padding_mode='median')(x)
+        assert with_median[0, 0, 2].tolist() == [1, 1]
+
+        # This is a special case: as we instantiated the tensor with integers,
+        # the mean (3/4) will be trucated to 0.
+        with_mean = tio.Pad(add_bottom_row, padding_mode='mean')(x)
+        assert with_mean[0, 0, 2].tolist() == [0, 0]
+        # So let's test with floats too
+        x = x.float()
+        with_mean = tio.Pad(add_bottom_row, padding_mode='mean')(x)
+        assert with_mean[0, 0, 2].tolist() == [0.75, 0.75]
+
+    def test_truncation_warning(self):
+        x = torch.ones(1, 1, 2, 2, dtype=torch.int)
+        pad = tio.Pad(1, padding_mode='mean')
+        with pytest.warns(RuntimeWarning):
+            pad(x)