torchio.transforms.augmentation.spatial.random_anisotropy

Complexity

Total Complexity

7

Size/Duplication

Total Lines	86
Duplicated Lines	0 %

Importance

Changes

0

4 Methods

Rating	Name	Duplication	Size	Complexity
A	RandomAnisotropy.apply_transform()	0	15	1
A	RandomAnisotropy.get_params()	0	8	1
A	RandomAnisotropy.__init__()	0	13	1
A	RandomAnisotropy.parse_axes()	0	8	4

from typing import Union, Tuple, Optional, List, Sequence
import torch
from ....torchio import TypeRangeFloat
from ....data.subject import Subject
from ....utils import to_tuple
from ... import SpatialTransform
from .. import RandomTransform
from ...preprocessing import Resample


class RandomAnisotropy(RandomTransform, SpatialTransform):
    r"""Downsample an image along an axis and upsample to initial space.

    This transform simulates an image that has been acquired using anisotropic
    spacing and resampled back to a more standard spacing.

    Args:
        axes: Axis or tuple of axes along which the image will be downsampled.
        downsampling: Downsampling factor :math:`m \gt 1`. If a tuple
            :math:`(a, b)` is provided then :math:`m \sim \mathcal{U}(a, b)`.
        image_interpolation: Image interpolation used to upsample the image back
            to its initial spacing. Downsampling is performed using nearest
            neighbor interpolation. See :ref:`Interpolation` for supported
            interpolation types.
        p: Probability that this transform will be applied.
        keys: See :py:class:`~torchio.transforms.Transform`.

    Example:
        >>> import torchio as tio
        >>> transform = tio.RandomAnisotropy(axes=1, downsampling=2.)   # Multiply spacing of second axis by 2
        >>> transform = tio.RandomAnisotropy(
        ...     axes=(0, 1, 2),
        ...     downsampling=(2, 5),
        ... )   # Multiply spacing of one of the 3 axes by a factor randomly chosen in [2, 5]
        >>> colin = tio.datasets.Colin27()
        >>> transformed = transform(colin)  # images have now anisotropic spacing
    """

    def __init__(
            self,
            axes: Union[int, Tuple[int, ...]] = (0, 1, 2),
            downsampling: TypeRangeFloat = (1.5, 5),
            image_interpolation: str = 'linear',
            p: float = 1,
            keys: Optional[Sequence[str]] = None,
            ):
        super().__init__(p=p, keys=keys)
        self.axes = self.parse_axes(axes)
        self.downsampling_range = self.parse_range(
            downsampling, 'downsampling', min_constraint=1)
        self.image_interpolation = self.parse_interpolation(image_interpolation)

    def get_params(
            self,
            axes: Tuple[int, ...],
            downsampling_range: Tuple[float, float],
            ) -> List[bool]:
        axis = axes[torch.randint(0, len(axes), (1,))]
        downsampling = self.sample_uniform(*downsampling_range).item()
        return axis, downsampling

    @staticmethod
    def parse_axes(axes: Union[int, Tuple[int, ...]]):
        axes_tuple = to_tuple(axes)
        for axis in axes_tuple:
            is_int = isinstance(axis, int)
            if not is_int or axis not in (0, 1, 2):
                raise ValueError('All axes must be 0, 1 or 2')
        return axes_tuple

    def apply_transform(self, subject: Subject) -> Subject:
        axis, downsampling = self.get_params(self.axes, self.downsampling_range)
        target_spacing = list(subject.spacing)
        target_spacing[axis] *= downsampling
        downsample = Resample(
            tuple(target_spacing),
            image_interpolation='nearest',
        )
        downsampled = downsample(subject)
        upsample = Resample(
            subject.get_first_image(),
            image_interpolation=self.image_interpolation,
        )
        upsampled = upsample(downsampled)
        return upsampled