Repository: TorchIO-project/torchio

Issue:
Pad to Target Shape
**🚀 Feature**
<!-- A clear and concise description of the feature proposal -->
A transform with most of the same functionality as `CropOrPad` but never crops, only pads to ensure a minimum target shape is achieved.

**Motivation**

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
Using a patch-based pipeline, I need to pad my images to ensure they are are not smaller than the patch size. For example, the S-I size varies between ~250 to >1000, and my patch size is usually 320 or 496. But I don't want to crop because I still want to be able to take multiple patches from the volumes.

**Pitch**
<!-- A clear and concise description of what you want to happen. -->

Either a new class to achieve this, built on top of the existing `CropOrPad` and just removing the crop section of `apply_transform`:

```python
class PadToTargetShape(tio.CropOrPad):
    """Pad, if necessary, to match a target shape."""
    def __init__(self, target_shape, **kwargs):
        super().__init__(target_shape=target_shape, **kwargs)

    def apply_transform(self, subject: tio.Subject) -> tio.Subject:
        subject.check_consistent_space()
        padding_params, _ = self.compute_crop_or_pad(subject)
        padding_kwargs = {'padding_mode': self.padding_mode}
        if padding_params is not None:
            pad = Pad(padding_params, **padding_kwargs)
            subject = pad(subject)  # type: ignore[assignment]
        return subject
```

Or change the existing CropOrPad class to add flags.

```python
class CropOrPad(SpatialTransform):
    def __init__(
        self,
        target_shape: Union[int, TypeTripletInt, None] = None,
        padding_mode: Union[str, float] = 0,
        mask_name: Optional[str] = None,
        labels: Optional[Sequence[int]] = None,
        apply_crop: bool = True,
        apply_pad: bool = True,
        **kwargs,
    ):
        .....
        self.apply_crop = apply_crop
        self.apply_pad = apply_pad

    def apply_transform(self, subject: Subject) -> Subject:
        subject.check_consistent_space()
        padding_params, cropping_params = self.compute_crop_or_pad(subject)
        padding_kwargs = {'padding_mode': self.padding_mode}
        if self.apply_pad and padding_params is not None:  # Added check here!
            pad = Pad(padding_params, **padding_kwargs)
            subject = pad(subject)  # type: ignore[assignment]
        if self.apply_crop and cropping_params is not None:  # Added check here!
            crop = Crop(cropping_params)
            subject = crop(subject)  # type: ignore[assignment]
        return subject
```

**Alternatives**

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
- This could be done as part of pre-processing, but it is less dynamic when wanting to experiment with patch sizes.
- It could be done by calculating the difference between the existing shape and target shape and using `tio.transforms.Pad`, but it would be much nicer just to have a Class already able to do this, along with the other functionality of `CropOrPad`, if required.

Maybe I'm missing something obvious, an easy way of achieving this already, if so let me know :)

Hints from issue comments:
@rickymwalsh: Seems like nobody else needs this so I guess I'll close for now, and at least the `PadToTargetShape` code is here for anyone who needs it.@fepegar: Apologies, @rickymwalsh. I actually think this makes sense and it would be great to have support. Modifying `CropOrPad` seems to be the simplest way forward. Would you like to contribute with a PR?

Solution patch:
diff --git a/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py b/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py
index c665f05..85c80d1 100644
--- a/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py
+++ b/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py
@@ -34,6 +34,10 @@ class CropOrPad(SpatialTransform):
             :attr:`mask_name`.
         labels: If a label map is used to generate the mask, sequence of labels
             to consider.
+        only_crop: If ``True``, padding will not be applied, only cropping will
+            be done. ``only_crop`` and ``only_pad`` cannot both be ``True``.
+        only_pad: If ``True``, cropping will not be applied, only padding will
+            be done. ``only_crop`` and ``only_pad`` cannot both be ``True``.
         **kwargs: See :class:`~torchio.transforms.Transform` for additional
             keyword arguments.
 
@@ -76,6 +80,8 @@ class CropOrPad(SpatialTransform):
         padding_mode: Union[str, float] = 0,
         mask_name: Optional[str] = None,
         labels: Optional[Sequence[int]] = None,
+        only_crop: bool = False,
+        only_pad: bool = False,
         **kwargs,
     ):
         if target_shape is None and mask_name is None:
@@ -111,6 +117,12 @@ class CropOrPad(SpatialTransform):
         self.mask_name = mask_name
         self.labels = labels
 
+        if only_pad and only_crop:
+            message = 'only_crop and only_pad cannot both be True'
+            raise ValueError(message)
+        self.only_crop = only_crop
+        self.only_pad = only_pad
+
     @staticmethod
     def _bbox_mask(mask_volume: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
         """Return 6 coordinates of a 3D bounding box from a given mask.
@@ -277,10 +289,10 @@ class CropOrPad(SpatialTransform):
         subject.check_consistent_space()
         padding_params, cropping_params = self.compute_crop_or_pad(subject)
         padding_kwargs = {'padding_mode': self.padding_mode}
-        if padding_params is not None:
+        if padding_params is not None and not self.only_crop:
             pad = Pad(padding_params, **self.get_base_args(), **padding_kwargs)
             subject = pad(subject)  # type: ignore[assignment]
-        if cropping_params is not None:
+        if cropping_params is not None and not self.only_pad:
             crop = Crop(cropping_params, **self.get_base_args())
             subject = crop(subject)  # type: ignore[assignment]
         return subject


Test patch:
diff --git a/tests/transforms/preprocessing/test_crop_pad.py b/tests/transforms/preprocessing/test_crop_pad.py
index ea70310..cb9cadb 100644
--- a/tests/transforms/preprocessing/test_crop_pad.py
+++ b/tests/transforms/preprocessing/test_crop_pad.py
@@ -212,3 +212,31 @@ class TestCropOrPad(TorchioTestCase):
             shape_a = crop(subject_a).image.shape
             shape_b = crop(subject_b).image.shape
             assert shape_a != shape_b
+
+    def test_only_crop_pad_true(self):
+        with pytest.raises(ValueError):
+            tio.CropOrPad((1, 2, 3), only_crop=True, only_pad=True)
+
+    def test_only_pad_true(self):
+        target_shape = 9, 21, 30
+        orig_shape = self.sample_subject['t1'].spatial_shape
+        expected_shape = tuple(
+            t if t > o else o for o, t in zip(orig_shape, target_shape)
+        )
+        transform = tio.CropOrPad(target_shape, only_pad=True)
+        transformed = transform(self.sample_subject)
+        for key in transformed:
+            result_shape = transformed[key].spatial_shape
+            assert result_shape == expected_shape
+
+    def test_only_crop_true(self):
+        target_shape = 9, 21, 30
+        orig_shape = self.sample_subject['t1'].spatial_shape
+        expected_shape = tuple(
+            t if t < o else o for o, t in zip(orig_shape, target_shape)
+        )
+        transform = tio.CropOrPad(target_shape, only_crop=True)
+        transformed = transform(self.sample_subject)
+        for key in transformed:
+            result_shape = transformed[key].spatial_shape
+            assert result_shape == expected_shape