Pad to Target Shape

[rickymwalsh]

🚀 Feature

A transform with most of the same functionality as CropOrPad but never crops, only pads to ensure a minimum target shape is achieved.

Motivation

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

Either a new class to achieve this, built on top of the existing CropOrPad and just removing the crop section of apply_transform:

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.

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

• 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 :)