Source code for torchio.transforms.preprocessing.spatial.pad

from numbers import Number
from typing import Union

import numpy as np
import nibabel as nib
import torch

from ....data.subject import Subject
from .bounds_transform import BoundsTransform, TypeBounds


[docs]class Pad(BoundsTransform): r"""Pad an image. Args: padding: Tuple :math:`(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})` defining the number of values padded to the edges of each axis. If the initial shape of the image is :math:`W \times H \times D`, the final shape will be :math:`(w_{ini} + W + w_{fin}) \times (h_{ini} + H + h_{fin}) \times (d_{ini} + D + d_{fin})`. If only three values :math:`(w, h, d)` are provided, then :math:`w_{ini} = w_{fin} = w`, :math:`h_{ini} = h_{fin} = h` and :math:`d_{ini} = d_{fin} = d`. If only one value :math:`n` is provided, then :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'``. **kwargs: See :class:`~torchio.transforms.Transform` for additional keyword arguments. .. _NumPy docs: https://numpy.org/doc/stable/reference/generated/numpy.pad.html """ # noqa: E501 PADDING_MODES = ( 'empty', 'edge', 'wrap', 'constant', 'linear_ramp', 'maximum', 'mean', 'median', 'minimum', 'reflect', 'symmetric', ) def __init__( self, padding: TypeBounds, padding_mode: Union[str, float] = 0, **kwargs ): super().__init__(padding, **kwargs) self.padding = padding self.check_padding_mode(padding_mode) self.padding_mode = padding_mode self.args_names = 'padding', 'padding_mode' @classmethod def check_padding_mode(cls, padding_mode): is_number = isinstance(padding_mode, Number) if not (padding_mode in cls.PADDING_MODES or is_number): message = ( f'Padding mode "{padding_mode}" not valid. Valid options are' f' {list(cls.PADDING_MODES)} or a number' ) raise KeyError(message) def apply_transform(self, subject: Subject) -> Subject: low = self.bounds_parameters[::2] for image in self.get_images(subject): new_origin = nib.affines.apply_affine(image.affine, -np.array(low)) new_affine = image.affine.copy() new_affine[:3, 3] = new_origin if isinstance(self.padding_mode, Number): kwargs = {'mode': 'constant', 'constant_values': self.padding_mode} else: kwargs = {'mode': self.padding_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) image.set_data(torch.as_tensor(padded)) image.affine = new_affine return subject def inverse(self): from .crop import Crop return Crop(self.padding)