Skip to content
Snippets Groups Projects
transforms.py 78.2 KiB
Newer Older
Lucia Hradecká's avatar
Lucia Hradecká committed
# ============================================================================================= #
#  Author:       Pavel Iakubovskii, ZFTurbo, ashawkey, Dominik Müller,                          #
#                Samuel Šuľan, Lucia Hradecká, Filip Lux                                        #
#  Copyright:    albumentations:    : https://github.com/albumentations-team                    #
#                Pavel Iakubovskii  : https://github.com/qubvel                                 #
#                ZFTurbo            : https://github.com/ZFTurbo                                #
#                ashawkey           : https://github.com/ashawkey                               #
#                Dominik Müller     : https://github.com/muellerdo                              #
#                Lucia Hradecká     : lucia.d.hradecka@gmail.com                                #
#                Filip Lux          : lux.filip@gmail.com                                       #
#                Samuel Šuľan                                                                   #
Lucia Hradecká's avatar
Lucia Hradecká committed
#                                                                                               #
#  Volumentations History:                                                                      #
#       - Original:                 https://github.com/albumentations-team/albumentations       #
#       - 3D Conversion:            https://github.com/ashawkey/volumentations                  #
#       - Continued Development:    https://github.com/ZFTurbo/volumentations                   #
#       - Enhancements:             https://github.com/qubvel/volumentations                    #
#       - Further Enhancements:     https://github.com/muellerdo/volumentations                 #
#       - Biomedical Enhancements:  https://gitlab.fi.muni.cz/cbia/bio-volumentations           #
#                                                                                               #
#  MIT License.                                                                                 #
#                                                                                               #
#  Permission is hereby granted, free of charge, to any person obtaining a copy                 #
#  of this software and associated documentation files (the "Software"), to deal                #
#  in the Software without restriction, including without limitation the rights                 #
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell                    #
#  copies of the Software, and to permit persons to whom the Software is                        #
#  furnished to do so, subject to the following conditions:                                     #
#                                                                                               #
#  The above copyright notice and this permission notice shall be included in all               #
#  copies or substantial portions of the Software.                                              #
#                                                                                               #
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR                   #
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,                     #
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE                  #
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER                       #
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,                #
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE                #
#  SOFTWARE.                                                                                    #
# ============================================================================================= #

import random
import numpy as np
from ..core.transforms_interface import DualTransform, ImageOnlyTransform
from ..augmentations import functional as F
from ..augmentations.spatial_funcional import get_affine_transform, parse_itk_interpolation
Lucia Hradecká's avatar
Lucia Hradecká committed
from ..random_utils import uniform, sample_range_uniform
from typing import List, Sequence, Tuple, Union
Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
from ..biovol_typing import *
from .utils import parse_limits, parse_coefs, parse_pads, to_tuple, validate_bbox, get_spatio_temporal_domain_limit,\
    to_spatio_temporal
Lucia Hradecká's avatar
Lucia Hradecká committed


# TODO anti_aliasing_downsample keep parameter or remove?
class Resize(DualTransform):
    """Resize input to the given shape.

        Internally, the ``skimage.transform.resize`` function is used.
        The ``interpolation``, ``border_mode``, ``ival``, ``mval``,
        and ``anti_aliasing_downsample`` arguments are forwarded to it. More details at:
        https://scikit-image.org/docs/stable/api/skimage.transform.html#skimage.transform.resize.

        Args:
            shape (tuple of ints): The desired image shape.

Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
                Must be ``(Z, Y, X)``.
Lucia Hradecká's avatar
Lucia Hradecká committed

Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
                The unspecified dimensions (C and T) are not affected.
Lucia Hradecká's avatar
Lucia Hradecká committed
            interpolation (int, optional): Order of spline interpolation.

                Defaults to ``1``.
            border_mode (str, optional): Values outside image domain are filled according to this mode.

                Defaults to ``'reflect'``.
            ival (float, optional): Value of `image` voxels outside of the `image` domain. Only applied when ``border_mode = 'constant'``.

                Defaults to ``0``.
            mval (float, optional): Value of `mask` and `float_mask` voxels outside of the domain. Only applied when ``border_mode = 'constant'``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            
                Defaults to ``0``.
            anti_aliasing_downsample (bool, optional): Controls if the Gaussian filter should be applied before
                downsampling. Recommended. 
                
                Defaults to ``True``.
            ignore_index (float | None, optional): If a float, then transformation of `mask` is done with 
                ``border_mode = 'constant'`` and ``mval = ignore_index``. 
                
                If ``None``, this argument is ignored.
                
                Defaults to ``None``.
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``1``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
    def __init__(self, shape: TypeSpatialShape, interpolation: int = 1, border_mode: str = 'reflect', ival: float = 0,
Lucia Hradecká's avatar
Lucia Hradecká committed
                 mval: float = 0, anti_aliasing_downsample: bool = True, ignore_index: Union[float, None] = None,
                 always_apply: bool = False, p: float = 1):
        
        super().__init__(always_apply, p)
        self.shape: TypeSpatioTemporalCoordinate = to_spatio_temporal(shape)
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.interpolation = interpolation
        self.border_mode = border_mode
        self.mask_mode = border_mode
        self.ival = ival
        self.mval = mval
        self.anti_aliasing_downsample = anti_aliasing_downsample
        if not (ignore_index is None):
            self.mask_mode = "constant"
            self.mval = ignore_index

    def apply(self, img, **params):
        return F.resize(img, input_new_shape=self.shape, interpolation=self.interpolation,
                        border_mode=self.border_mode, cval=self.ival,
                        anti_aliasing_downsample=self.anti_aliasing_downsample)

    def apply_to_mask(self, mask, **params):
        return F.resize(mask, input_new_shape=self.shape, interpolation=0,
                        border_mode=self.mask_mode, cval=self.mval, anti_aliasing_downsample=False,
                        mask=True)

    def apply_to_float_mask(self, mask, **params):
        return F.resize(mask, input_new_shape=self.shape, interpolation=self.interpolation,
                        border_mode=self.mask_mode, cval=self.mval, anti_aliasing_downsample=False,
                        mask=True)

    def apply_to_keypoints(self, keypoints, **params):
        return F.resize_keypoints(keypoints,
                                  domain_limit=params['domain_limit'],
                                  new_shape=self.shape)

    """
    def apply_to_bboxes(self, bboxes, **params):
        for bbox in bboxes:
            new_bbox = F.resize_keypoints(bbox,
                                          input_new_shape=self.shape,
                                          original_shape=params['original_shape'],
                                          keep_all=True)

            if validate_bbox(bbox, new_bbox, min_overlay_ratio):
                res.append(new_bbox)

        return res
    """

    def get_params(self, **data):

        # read shape of the original image
        domain_limit: TypeSpatioTemporalCoordinate = get_spatio_temporal_domain_limit(data)

        return {
            "domain_limit": domain_limit,
        }
Lucia Hradecká's avatar
Lucia Hradecká committed
        
    def __repr__(self):
        return f'Resize({self.shape}, {self.interpolation}, {self.border_mode} , {self.ival}, {self.mval},' \
               f'{self.anti_aliasing_downsample}, {self.always_apply}, {self.p})'
Lucia Hradecká's avatar
Lucia Hradecká committed


class Scale(DualTransform):
    """Rescale the input image content by the given scale. The image shape remains unchanged.
Lucia Hradecká's avatar
Lucia Hradecká committed

        Args:
            scales (float|List[float], optional): Value by which the input should be scaled.

                Must be either of: ``S``, ``[S_Z, S_Y, S_X]``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a float, then all spatial dimensions are scaled by it (equivalent to ``[S, S, S]``).

                The unspecified dimensions (C and T) are not affected.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``1``.
            interpolation (str, optional): SimpleITK interpolation type for `image` and `float_mask`.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Must be one of ``linear``, ``nearest``, ``bspline``, ``gaussian``.

                For `mask`, the ``nearest`` interpolation is always used.

                Defaults to ``linear``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            spacing (float | Tuple[float, float, float] | None, optional): Voxel spacing for individual spatial dimensions.

                Must be either of: ``S``, ``(S1, S2, S3)``, or ``None``.

                If ``None``, equivalent to ``(1, 1, 1)``.

                If a float ``S``, equivalent to ``(S, S, S)``.

                Otherwise, a scale for each spatial dimension must be given.

                Defaults to ``None``.
            border_mode (str, optional): Values outside image domain are filled according to this mode.

                Defaults to ``'constant'``.
            ival (float, optional): Value of `image` voxels outside of the `image` domain. Only applied when ``border_mode = 'constant'``.

                Defaults to ``0``.
            mval (float, optional): Value of `mask` and `float_mask` voxels outside of the domain. Only applied when ``border_mode = 'constant'``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``0``.
            ignore_index (float | None, optional): If a float, then transformation of `mask` is done with 
                ``border_mode = 'constant'`` and ``mval = ignore_index``. 
                
                If ``None``, this argument is ignored.

                Defaults to ``None``.
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``1``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
    def __init__(self, scales: Union[float, TypeTripletFloat] = 1,
                 interpolation: str = 'linear', spacing: Union[float, TypeTripletFloat] = None,
Lucia Hradecká's avatar
Lucia Hradecká committed
                 border_mode: str = 'constant', ival: float = 0, mval: float = 0,
                 ignore_index: Union[float, None] = None, always_apply: bool = False, p: float = 1):
        super().__init__(always_apply, p)
        self.scale = parse_coefs(scales, identity_element=1.)
        self.interpolation: str = parse_itk_interpolation(interpolation)
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.spacing: TypeTripletFloat = parse_coefs(spacing, identity_element=1.)
        self.border_mode = border_mode              # not implemented
        self.mask_mode = border_mode                # not implemented
        self.ival = ival
        self.mval = mval
        if not (ignore_index is None):
            self.mask_mode = "constant"
            self.mval = ignore_index

    def apply(self, img, **params):
        return F.affine(img,
                        scales=self.scale,
                        interpolation=self.interpolation,
                        border_mode=self.border_mode,
                        value=self.ival,
                        spacing=self.spacing)

    def apply_to_mask(self, mask, **params):
        interpolation = parse_itk_interpolation('nearest')   # refers to 'sitkNearestNeighbor'
Filip Lux's avatar
Filip Lux committed
        return F.affine(np.expand_dims(mask, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                        scales=self.scale,
                        interpolation=interpolation,
                        border_mode=self.mask_mode,
                        value=self.mval,
Filip Lux's avatar
Filip Lux committed
                        spacing=self.spacing)[0]
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_float_mask(self, mask, **params):
Filip Lux's avatar
Filip Lux committed
        return F.affine(np.expand_dims(mask, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                        scales=self.scale,
                        interpolation=self.interpolation,
                        border_mode=self.mask_mode,
                        value=self.mval,
Filip Lux's avatar
Filip Lux committed
                        spacing=self.spacing)[0]
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_keypoints(self, keypoints, **params):
        return F.affine_keypoints(keypoints,
                                  scales=self.scale,
                                  spacing = self.spacing,
                                  domain_limit=params['domain_limit'])

    """
    def apply_to_bboxes(self, bboxes, **params):
        for bbox in bboxes:
            new_bbox = F.affine_keypoints(bbox,
                                          scales=self.scale,
                                          domain_limit=params['domain_limit'],
                                          spacing = self.spacing,
                                          keep_all=True)

            if validate_bbox(bbox, new_bbox):
                res.append(new_bbox)

        return res
    """

    def get_params(self, **data):
        domain_limit: TypeSpatioTemporalCoordinate = get_spatio_temporal_domain_limit(data)
        return {'domain_limit': domain_limit}

Lucia Hradecká's avatar
Lucia Hradecká committed
    def __repr__(self):
        return f'Scale({self.scale}, {self.interpolation}, {self.border_mode}, {self.ival}, {self.mval},' \
               f'{self.always_apply}, {self.p})'


class RandomScale(DualTransform):
    """Randomly rescale the input image content by the given scale. The image shape remains unchanged.
Lucia Hradecká's avatar
Lucia Hradecká committed

        Args:
            scaling_limit (float | Tuple[float], optional): Limits of scaling factors.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Must be either of: ``S``, ``(S1, S2)``, ``(S_Z, S_Y, S_X)``, or ``(S_Z1, S_Z2, S_Y1, S_Y2, S_X1, S_X2)``.

                If a float ``S``, then all spatial dimensions are scaled by a random number drawn uniformly from
                the interval [1/S, S] (equivalent to inputting ``(1/S, S, 1/S, S, 1/S, S)``).
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a tuple of 2 floats, then all spatial dimensions are scaled by a random number drawn uniformly
                from the interval [S1, S2] (equivalent to inputting ``(S1, S2, S1, S2, S1, S2)``).

                If a tuple of 3 floats, then an interval [1/S_a, S_a] is constructed for each spatial
Lucia Hradecká's avatar
Lucia Hradecká committed
                dimension and the scale is randomly drawn from it
                (equivalent to inputting ``(1/S_Z, S_Z, 1/S_Y, S_Y, 1/S_X, S_X)``).
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a tuple of 6 floats, the scales for individual spatial dimensions are randomly drawn from the
                respective intervals [S_Z1, S_Z2], [S_Y1, S_Y2], [S_X1, S_X2].

                The unspecified dimensions (C and T) are not affected.

Lucia Hradecká's avatar
Lucia Hradecká committed

            interpolation (str, optional): SimpleITK interpolation type for `image` and `float_mask`.

                Must be one of ``linear``, ``nearest``, ``bspline``, ``gaussian``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                For `mask`, the ``nearest`` interpolation is always used.

                Defaults to ``linear``.
Lucia Hradecká's avatar
Lucia Hradecká committed

            spacing (float | Tuple[float, float, float] | None, optional): Voxel spacing for individual spatial dimensions.

                Must be either of: ``S``, ``(S1, S2, S3)``, or ``None``.

                If ``None``, equivalent to ``(1, 1, 1)``.

                If a float ``S``, equivalent to ``(S, S, S)``.

                Otherwise, a scale for each spatial dimension must be given.

                Defaults to ``None``.

            border_mode (str, optional): Values outside image domain are filled according to the mode.

                Defaults to ``'constant'``.

            ival (float, optional): Value of `image` voxels outside of the `image` domain. Only applied when ``border_mode = 'constant'``.

                Defaults to ``0``.

            mval (float, optional): Value of `mask` and `float_mask` voxels outside of the domain. Only applied when ``border_mode = 'constant'``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``0``.

            ignore_index (float | None, optional): If a float, then transformation of `mask` is done with 
                ``border_mode = 'constant'`` and ``mval = ignore_index``. 
                
                If ``None``, this argument is ignored.

                Defaults to ``None``.

            always_apply (bool, optional): Always apply this transformation in composition.

                Defaults to ``False``.

            p (float, optional): Chance of applying this transformation in composition.

                Defaults to ``0.5``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """      
    def __init__(self, scaling_limit: Union[float, TypePairFloat, TypeTripletFloat, TypeSextetFloat] = (0.9, 1.1),
                 interpolation: str = 'linear', spacing: Union[float, TypeTripletFloat] = None,
Lucia Hradecká's avatar
Lucia Hradecká committed
                 border_mode: str = 'constant', ival: float = 0, mval: float = 0,
                 ignore_index: Union[float, None] = None, always_apply: bool = False, p: float = 0.5):
        super().__init__(always_apply, p)
        self.scaling_limit: TypeSextetFloat = parse_limits(scaling_limit, scale=True)
        self.interpolation: str = parse_itk_interpolation(interpolation)
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.spacing: TypeTripletFloat = parse_coefs(spacing, identity_element=1.)
        self.border_mode = border_mode
        self.mask_mode = border_mode
        self.ival: float = ival
        self.mval: float = mval
        if not (ignore_index is None):
            self.mask_mode = "constant"
            self.mval = ignore_index

    def get_params(self, **data):
        # set parameters of the transform
        domain_limit: TypeSpatioTemporalCoordinate = get_spatio_temporal_domain_limit(data)
Lucia Hradecká's avatar
Lucia Hradecká committed
        scale = sample_range_uniform(self.scaling_limit)

        return {
            "domain_limit": domain_limit,
Lucia Hradecká's avatar
Lucia Hradecká committed
            "scale": scale,
        }

    def apply(self, img, **params):
        return F.affine(img,
                        scales=params["scale"],
                        interpolation=self.interpolation,
                        border_mode=self.border_mode,
                        value=self.ival,
                        spacing=self.spacing)

    def apply_to_mask(self, mask, **params):
        interpolation = parse_itk_interpolation('nearest')   # refers to 'sitkNearestNeighbor'
Filip Lux's avatar
Filip Lux committed
        return F.affine(np.expand_dims(mask, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                        scales=params["scale"],
                        interpolation=interpolation,
                        border_mode=self.mask_mode,
                        value=self.mval,
Filip Lux's avatar
Filip Lux committed
                        spacing=self.spacing)[0]
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_float_mask(self, mask, **params):
Filip Lux's avatar
Filip Lux committed
        return F.affine(np.expand_dims(mask, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                        scales=params["scale"],
                        interpolation=self.interpolation,
                        border_mode=self.mask_mode,
                        value=self.mval,
Filip Lux's avatar
Filip Lux committed
                        spacing=self.spacing)[0]
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_keypoints(self, keypoints, **params):
        return F.affine_keypoints(keypoints,
                                  scales=params["scale"],
                                  spacing=self.spacing,
                                  domain_limit=params['domain_limit'])

Lucia Hradecká's avatar
Lucia Hradecká committed
    def __repr__(self):
        return f'RandomScale({self.scaling_limit}, {self.interpolation}, {self.always_apply}, {self.p})'


class RandomRotate90(DualTransform):
    """Rotation of input by 0, 90, 180, or 270 degrees around the specified spatial axes.

        Args:
            axes (List[int], optional): List of axes around which the input is rotated. Recognised axis symbols are
                ``1`` for Z, ``2`` for Y, and ``3`` for X. A single axis can occur multiple times in the list.
                If ``shuffle_axis = False``, the order of axes determines the order of transformations.
                If ``None``, will be rotated around all spatial axes.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``None``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            shuffle_axis (bool, optional): If set to ``True``, the order of rotations is random.

                Defaults to ``False``.
            factor (int, optional): Number of times the array is rotated by 90 degrees. If ``None``, will be chosen randomly.

                Defaults to ``None``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``0.5``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
    def __init__(self, axes: List[int] = None, shuffle_axis: bool = False, factor=None,
Lucia Hradecká's avatar
Lucia Hradecká committed
                 always_apply: bool = False, p: float = 0.5):
        super().__init__(always_apply, p)
        self.axes = axes
        self.shuffle_axis = shuffle_axis
        self.factor = factor
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply(self, img, **params):
        for factor, axes in zip(params["factor"], params["rotation_around"]):
            img = np.rot90(img, factor, axes=axes)
        return img

    def apply_to_mask(self, mask, **params):
Filip Lux's avatar
Filip Lux committed
        for rot, factor in zip(params["rotation_around"], params["factor"]):
            mask = np.rot90(mask, factor, axes=(rot[0] - 1, rot[1] - 1))
Lucia Hradecká's avatar
Lucia Hradecká committed
        return mask

Filip Lux's avatar
Filip Lux committed
    def apply_to_keypoints(self, keypoints, **params):
Filip Lux's avatar
Filip Lux committed
        for rot, factor in zip(params["rotation_around"], params["factor"]):
            keypoints = F.rot90_keypoints(keypoints,
                                          factor=factor,
                                          axes=(rot[0], rot[1]),
Filip Lux's avatar
Filip Lux committed
                                          img_shape=params['img_shape'])
        return keypoints

Lucia Hradecká's avatar
Lucia Hradecká committed
    def get_params(self, **data):

        # Rotate by all axis by default
        if self.axes is None:
            self.axes = [1, 2, 3]

        # Create all combinations for rotating
        axes_to_rotate = {1: (2, 3), 2: (1, 3), 3: (1, 2)}
        rotation_around = []
        for i in self.axes:
            if i in axes_to_rotate.keys():
                rotation_around.append(axes_to_rotate[i])

        # shuffle order of axis
        if self.shuffle_axis:
            random.shuffle(rotation_around)

        # choose angle to rotate
        if self.factor is None:
            factor = [random.randint(0, 3) for _ in range(len(rotation_around))]
        else:
            factor = [self.factor]
            rotation_around = [(1, 2)]
            print('ROT90', factor, rotation_around)

Filip Lux's avatar
Filip Lux committed
        img_shape = np.array(data['image'].shape[1:4])

Lucia Hradecká's avatar
Lucia Hradecká committed
        return {"factor": factor,
Filip Lux's avatar
Filip Lux committed
                "rotation_around": rotation_around,
                "img_shape": img_shape}
Lucia Hradecká's avatar
Lucia Hradecká committed

    def __repr__(self):
        return f'RandomRotate90({self.axes}, {self.always_apply}, {self.p})'


class Flip(DualTransform):
    """Flip input around the specified spatial axes.

        Args:
            axes (List[int], optional): List of axes around which is flip done. Recognised axis symbols are
                ``1`` for Z, ``2`` for Y, and ``3`` for X. If ``None``, will be flipped around all spatial axes.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``None``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``1``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
    def __init__(self, axes: List[int] = None, always_apply=False, p=1):
        super().__init__(always_apply, p)
        self.axes = axes

    def apply(self, img, **params):
        return np.flip(img, params["axes"])

    def apply_to_mask(self, mask, **params):
        # Mask has no dimension channel
        return np.flip(mask, axis=[item - 1 for item in params["axes"]])

Filip Lux's avatar
Filip Lux committed
    def apply_to_keypoints(self, keypoints, **params):
Filip Lux's avatar
Filip Lux committed
        return F.flip_keypoints(keypoints,
                                axes=params['axes'],
                                img_shape=params['img_shape'])

Lucia Hradecká's avatar
Lucia Hradecká committed
    def get_params(self, **data):
Filip Lux's avatar
Filip Lux committed
        axes = [1, 2, 3] if self.axes is None else self.axes
        img_shape = np.array(data['image'].shape[1:4])
        return {"axes": axes,
                "img_shape": img_shape}
Lucia Hradecká's avatar
Lucia Hradecká committed

    def __repr__(self):
        return f'Flip({self.axes}, {self.always_apply}, {self.p})'


class RandomFlip(DualTransform):
    """Flip input around a set of axes randomly chosen from the input list of axis combinations.

        Args:
Filip Lux's avatar
Filip Lux committed
            axes_to_choose (List[Tuple[int]] or None, optional): List of axis indices from which one option
Lucia Hradecká's avatar
Lucia Hradecká committed
                is randomly chosen. Recognised axis symbols are ``1`` for Z, ``2`` for Y, and ``3`` for X.
                The image will be flipped around all axes in the chosen combination.

                If ``None``, a random subset of spatial axes is chosen, corresponding to inputting
Filip Lux's avatar
Filip Lux committed
                ``[(,), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``None``.
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``0.5``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
    def __init__(self, axes_to_choose: Union[None, List[Tuple[int]]] = None, always_apply=False, p=0.5):
        super().__init__(always_apply, p)
Filip Lux's avatar
Filip Lux committed

        # TODO: check if input value `axes_to_choose` valid
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.axes = axes_to_choose

    def apply(self, img, **params):
        return np.flip(img, params["axes"])

    def apply_to_mask(self, mask, **params):
        # Mask has no dimension channel
        return np.flip(mask, axis=[item - 1 for item in params["axes"]])

Filip Lux's avatar
Filip Lux committed
    def apply_to_keypoints(self, keypoints, keep_all=False, **params):
        return F.flip_keypoints(keypoints,
                                axes=params['axes'],
                                img_shape=params['img_shape'])

Lucia Hradecká's avatar
Lucia Hradecká committed
    def get_params(self, **data):
        
Filip Lux's avatar
Filip Lux committed
        to_choose = [1, 2, 3] if self.axes is None else self.axes
        axes = random.sample(to_choose, random.randint(0, len(to_choose)))
Filip Lux's avatar
Filip Lux committed
        img_shape = np.array(data['image'].shape[1:4])
        return {"axes": axes,
                "img_shape": img_shape}
Lucia Hradecká's avatar
Lucia Hradecká committed

    def __repr__(self):
        return f'RandomFlip({self.axes}, {self.always_apply}, {self.p})'
Lucia Hradecká's avatar
Lucia Hradecká committed


class CenterCrop(DualTransform):
    """Crops the central region of the input of given size.
          
        Unlike ``CenterCrop`` from `Albumentations`, this transform pads the input in dimensions
        where the input is smaller than the ``shape`` with ``numpy.pad``. The ``border_mode``, ``ival`` and ``mval``
        arguments are forwarded to ``numpy.pad`` if padding is necessary. More details at:
        https://numpy.org/doc/stable/reference/generated/numpy.pad.html.

        Args:
            shape (Tuple[int]): The desired shape of input.

                Must be ``[Z, Y, X]``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            border_mode (str, optional): Values outside image domain are filled according to this mode.

                Defaults to ``'reflect'``.
            ival (float | Sequence, optional): Values of `image` voxels outside of the `image` domain.
                Only applied when ``border_mode = 'constant'`` or ``border_mode = 'linear_ramp'``.

                Defaults to ``(0, 0)``.
            mval (float | Sequence, optional): Values of `mask` voxels outside of the `mask` domain.
                Only applied when ``border_mode = 'constant'`` or ``border_mode = 'linear_ramp'``.

                Defaults to ``(0, 0)``.
            ignore_index (float | None, optional): If a float, then transformation of `mask` is done with 
                ``border_mode = 'constant'`` and ``mval = ignore_index``. 
                
                If ``None``, this argument is ignored.

                Defaults to ``None``.
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``1``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
    def __init__(self, shape: TypeSpatialShape, border_mode: str = "reflect", ival: Union[Sequence[float], float] = (0, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                 mval: Union[Sequence[float], float] = (0, 0), ignore_index: Union[float, None] = None,
                 always_apply: bool = False, p: float = 1.0):
        super().__init__(always_apply, p)
Filip Lux's avatar
Filip Lux committed
        self.output_shape = np.asarray(shape, dtype=np.intc)  # TODO: make it len 3
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.border_mode = border_mode
        self.mask_mode = border_mode
        self.ival = ival
        self.mval = mval
        
        if not (ignore_index is None):
            self.mask_mode = "constant"
            self.mval = ignore_index

    def apply(self, img, **params):
Filip Lux's avatar
Filip Lux committed
        return F.crop(img,
                      crop_shape=self.output_shape,
                      crop_position=params['crop_position'],
                      pad_dims=params['pad_dims'],
                      border_mode=self.mask_mode, cval=self.mval, mask=False)
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_mask(self, mask, **params):
Filip Lux's avatar
Filip Lux committed
        return F.crop(mask,
                      crop_shape=self.output_shape,
                      crop_position=params['crop_position'],
                      pad_dims=params['pad_dims'],
                      border_mode=self.mask_mode, cval=self.mval, mask=True)

Filip Lux's avatar
Filip Lux committed
    def apply_to_keypoints(self, keypoints, keep_all=False, **params):
        return F.crop_keypoints(keypoints,
                                crop_shape=self.output_shape,
                                crop_position=params['crop_position'],
                                pad_dims=params['pad_dims'],
                                keep_all=keep_all)

Filip Lux's avatar
Filip Lux committed
    def get_params(self, **data):
        # get crop coordinates, position of the corner closest to the image origin
        img_spatial_shape = np.array(data['image'].shape[1:4])
        position: TypeSpatialCoordinate = (img_spatial_shape - self.output_shape) // 2
        position = np.maximum(position, 0).astype(int)
        pad_dims = F.get_pad_dims(img_spatial_shape, self.output_shape)

        return {'crop_position': position,
                'pad_dims': pad_dims}
Lucia Hradecká's avatar
Lucia Hradecká committed

    def __repr__(self):
Filip Lux's avatar
Filip Lux committed
        return f'CenterCrop({self.output_shape}, {self.always_apply}, {self.p})'
Lucia Hradecká's avatar
Lucia Hradecká committed


class RandomCrop(DualTransform):
    """Randomly crops a region of given size from the input.

        Unlike ``RandomCrop`` from `Albumentations`, this transform pads the input in dimensions
        where the input is smaller than the ``shape`` with ``numpy.pad``. The ``border_mode``, ``ival`` and ``mval``
        arguments are forwarded to ``numpy.pad`` if padding is necessary. More details at:
        https://numpy.org/doc/stable/reference/generated/numpy.pad.html.

        Args:
            shape (Tuple[int]): The desired shape of input.

                Must be ``[Z, Y, X]``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            border_mode (str, optional): Values outside image domain are filled according to this mode.

                Defaults to ``'reflect'``.
            ival (float | Sequence, optional): Values of `image` voxels outside of the `image` domain.
                Only applied when ``border_mode = 'constant'`` or ``border_mode = 'linear_ramp'``.

                Defaults to ``(0, 0)``.
            mval (float | Sequence, optional): Values of `mask` voxels outside of the `mask` domain.
                Only applied when ``border_mode = 'constant'`` or ``border_mode = 'linear_ramp'``.

                Defaults to ``(0, 0)``.
            ignore_index (float | None, optional): If a float, then transformation of `mask` is done with 
                ``border_mode = 'constant'`` and ``mval = ignore_index``. 
                
                If ``None``, this argument is ignored.

                Defaults to ``None``.
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``1``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
    def __init__(self, shape: TypeSpatialShape, border_mode: str = "reflect", ival: Union[Sequence[float], float] = (0, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                 mval: Union[Sequence[float], float] = (0, 0), ignore_index: Union[float, None] = None,
                 always_apply: bool = False, p: float = 1.0):
        super().__init__(always_apply, p)
Filip Lux's avatar
Filip Lux committed
        self.output_shape = np.asarray(shape, dtype=np.intc)
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.border_mode = border_mode
        self.mask_mode = border_mode
        self.ival = ival
        self.mval = mval

        if not (ignore_index is None):
            self.mask_mode = "constant"
            self.mval = ignore_index

Filip Lux's avatar
Filip Lux committed
    def apply(self, img, **params):
        return F.crop(img,
                      crop_shape=self.output_shape,
                      crop_position=params['crop_position'],
                      pad_dims=params['pad_dims'],
                      border_mode=self.mask_mode, cval=self.mval, mask=False)
Lucia Hradecká's avatar
Lucia Hradecká committed

Filip Lux's avatar
Filip Lux committed
    def apply_to_mask(self, mask, **params):
        return F.crop(mask,
                      crop_shape=self.output_shape,
                      crop_position=params['crop_position'],
                      pad_dims=params['pad_dims'],
                      border_mode=self.mask_mode, cval=self.mval, mask=True)
Filip Lux's avatar
Filip Lux committed
    def apply_to_keypoints(self, keypoints, keep_all=False, **params):
        return F.crop_keypoints(keypoints,
                                crop_shape=self.output_shape,
                                crop_position=params['crop_position'],
                                pad_dims=params['pad_dims'],
                                keep_all=keep_all)

Lucia Hradecká's avatar
Lucia Hradecká committed
    def get_params(self, **data):
Filip Lux's avatar
Filip Lux committed
        # get crop coordinates, position of the corner closest to the image origin
        img_spatial_shape = np.array(data['image'].shape[1:4])
        ranges: TypeSpatialShape = np.maximum(img_spatial_shape - self.output_shape, 0)
        position = np.array([random.randint(0, r) for r in ranges])
        pad_dims = F.get_pad_dims(img_spatial_shape, self.output_shape)
        return {'crop_position': position,
                'pad_dims': pad_dims}
Lucia Hradecká's avatar
Lucia Hradecká committed

    def __repr__(self):
Filip Lux's avatar
Filip Lux committed
        return f'RandomCrop({self.output_shape}, {self.always_apply}, {self.p})'
Lucia Hradecká's avatar
Lucia Hradecká committed


class RandomAffineTransform(DualTransform):
Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
    """Affine transformation of the input image with randomly chosen parameters. Image shape remains unchanged.
Lucia Hradecká's avatar
Lucia Hradecká committed

        Args:
            angle_limit (Tuple[float] | float, optional): Intervals in degrees from which angles of
                rotation for the spatial axes are chosen.

Filip Lux's avatar
Filip Lux committed
                Must be either of: ``A``, ``(A1, A2)``, ``(A1, A2, A3)``, or ``(A_Z1, A_Z2, A_Y1, A_Y2, A_X1, A_X2)``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a float, equivalent to ``(-A, A, -A, A, -A, A)``.

                If a tuple with 2 items, equivalent to ``(A1, A2, A1, A2, A1, A2)``.

Filip Lux's avatar
Filip Lux committed
                If a tuple with 3 items, equivalent to ``(-A1, A1, -A2, A2, -A3, A3)``.

Lucia Hradecká's avatar
Lucia Hradecká committed
                If a tuple with 6 items, angle of rotation is randomly chosen from an interval [A_a1, A_a2] for each
                spatial axis.

                Defaults to ``(15, 15, 15)``.
Filip Lux's avatar
Filip Lux committed
            translation_limit (Tuple[float] | float | None, optional): Intervals from which the translation parameters
Lucia Hradecká's avatar
Lucia Hradecká committed
                for the spatial axes are chosen.

Filip Lux's avatar
Filip Lux committed
                Must be either of: ``T``, ``(T1, T2)``, ``(T1, T2, T3)``, or ``(T_Z1, T_Z2, T_Y1, T_Y2, T_X1, T_X2)``.
Lucia Hradecká's avatar
Lucia Hradecká committed

Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
                If a float, equivalent to ``(2-T, T, 2-T, T, 2-T, T)``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a tuple with 2 items, equivalent to ``(T1, T2, T1, T2, T1, T2)``.

Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
                If a tuple with 3 items, equivalent to ``(2-T1, T1, 2-T2, T2, 2-T3, T3)``.
Lucia Hradecká's avatar
Lucia Hradecká committed
                If a tuple with 6 items, the translation parameter is randomly chosen from an interval [T_a1, T_a2] for
                each spatial axis.

                Defaults to ``(0, 0, 0)``.
            scaling_limit (Tuple[float] | float, optional): Intervals from which the scales for the spatial axes are chosen.

Filip Lux's avatar
Filip Lux committed
                Must be either of: ``S``, ``(S1, S2)``, ``(S1, S2, S3)``, or ``(S_Z1, S_Z2, S_Y1, S_Y2, S_X1, S_X2)``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a float, equivalent to ``(1/S, S, 1/S, S, 1/S, S)``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                If a tuple with 2 items, equivalent to ``(S1, S2, S1, S2, S1, S2)``.

                If a tuple with 3 items, equivalent to ``(1/S1, S1, 1/S2, S2, 1/S3, S3)``.
Lucia Hradecká's avatar
Lucia Hradecká committed
                If a tuple with 6 items, the scale is randomly chosen from an interval [S_a1, S_a2] for
                each spatial axis.

                Defaults to ``(1., 1., 1.)``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            spacing (float | Tuple[float, float, float] | None, optional): Voxel spacing for individual spatial dimensions.

                Must be either of: ``S``, ``(S1, S2, S3)``, or ``None``.

                If ``None``, equivalent to ``(1, 1, 1)``.

                If a float ``S``, equivalent to ``(S, S, S)``.

                Otherwise, a scale for each spatial dimension must be given.

                Defaults to ``None``.
            change_to_isotropic (bool, optional): Change data from anisotropic to isotropic.

                Defaults to ``False``.
            interpolation (str, optional): SimpleITK interpolation type for `image` and `float_mask`.

                Must be one of ``linear``, ``nearest``, ``bspline``, ``gaussian``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                For `mask`, the ``nearest`` interpolation is always used.

                Defaults to ``linear``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            border_mode (str, optional): Values outside image domain are filled according to this mode.

                Defaults to ``'constant'``.
            ival (float, optional): Value of `image` voxels outside of the `image` domain. Only applied when ``border_mode = 'constant'``.

                Defaults to ``0``.
            mval (float, optional): Value of `mask` and `float_mask` voxels outside of the domain. Only applied when ``border_mode = 'constant'``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``0``.
            ignore_index (float | None, optional): If a float, then transformation of `mask` is done with 
                ``border_mode = 'constant'`` and ``mval = ignore_index``. 
                
                If ``None``, this argument is ignored.

                Defaults to ``None``.
            always_apply (bool, optional): Always apply this transformation in composition. 
            
                Defaults to ``False``.
            p (float, optional): Chance of applying this transformation in composition. 
            
                Defaults to ``0.5``.

        Targets:
            image, mask, float mask, key points, bounding boxes
Lucia Hradecká's avatar
Lucia Hradecká committed
    """
Filip Lux's avatar
Filip Lux committed
    def __init__(self, angle_limit: Union[float, TypePairFloat, TypeTripletFloat, TypeSextetFloat] = (15., 15., 15.),
                 translation_limit: Union[float, TypePairFloat, TypeTripletFloat, TypeSextetFloat] = (0., 0., 0.),
                 scaling_limit: Union[float, TypePairFloat, TypeTripletFloat, TypeSextetFloat] = (1., 1., 1.),
Lucia Hradecká's avatar
Lucia Hradecká committed
                 spacing: Union[float, TypeTripletFloat] = None,
                 change_to_isotropic: bool = False,
                 interpolation: str = 'linear',
Lucia Hradecká's avatar
Lucia Hradecká committed
                 border_mode: str = 'constant', ival: float = 0, mval: float = 0,
                 ignore_index: Union[float, None] = None, always_apply: bool = False, p: float = 0.5):
        super().__init__(always_apply, p)
        self.angle_limit: TypeSextetFloat = parse_limits(angle_limit)
        self.translation_limit: TypeSextetFloat = parse_limits(translation_limit)
        self.scaling_limit: TypeSextetFloat = parse_limits(scaling_limit, scale=True)
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.spacing: TypeTripletFloat = parse_coefs(spacing, identity_element=1)
        self.interpolation: int = parse_itk_interpolation(interpolation)
Lucia Hradecká's avatar
Lucia Hradecká committed
        self.border_mode = border_mode                 # not used
        self.mask_mode = border_mode                   # not used
        self.ival = ival
        self.mval = mval
        self.keep_scale = not change_to_isotropic

        if ignore_index is not None:
            self.mask_mode = "constant"
            self.mval = ignore_index

    def apply(self, img, **params):
        return F.affine(img,
                        scales=params["scale"],
                        degrees=params["angles"],
                        translation=params["translation"],
                        interpolation=self.interpolation,
                        border_mode=self.border_mode,
                        value=self.ival,
                        spacing=self.spacing)

    def apply_to_mask(self, mask, **params):
        interpolation = parse_itk_interpolation('nearest')   # refers to 'sitkNearestNeighbor'
Filip Lux's avatar
Filip Lux committed
        return F.affine(np.expand_dims(mask, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                        scales=params["scale"],
                        degrees=params["angles"],
                        translation=params["translation"],
                        interpolation=interpolation,
                        border_mode=self.mask_mode,
                        value=self.mval,
Filip Lux's avatar
Filip Lux committed
                        spacing=self.spacing)[0]
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_float_mask(self, mask, **params):
Filip Lux's avatar
Filip Lux committed
        return F.affine(np.expand_dims(mask, 0),
Lucia Hradecká's avatar
Lucia Hradecká committed
                        scales=params["scale"],
                        degrees=params["angles"],
                        translation=params["translation"],
                        interpolation=self.interpolation,
                        border_mode=self.mask_mode,
                        value=self.mval,
Filip Lux's avatar
Filip Lux committed
                        spacing=self.spacing)[0]
Lucia Hradecká's avatar
Lucia Hradecká committed

    def apply_to_keypoints(self, keypoints, **params):
        return F.affine_keypoints(keypoints,
                                  scales=params["scale"],
                                  degrees=params["angles"],
                                  translation=params["translation"],
                                  spacing=self.spacing,
                                  domain_limit=params['domain_limit'])

Lucia Hradecká's avatar
Lucia Hradecká committed
    def get_params(self, **data):

        # set parameters of the transform
        scales = sample_range_uniform(self.scaling_limit)
        angles = sample_range_uniform(self.angle_limit)
        translation = sample_range_uniform(self.translation_limit)
        domain_limit = get_spatio_temporal_domain_limit(data)
Lucia Hradecká's avatar
Lucia Hradecká committed

        return {
            "scale": scales,
            "angles": angles,
            "translation": translation,
            "domain_limit": domain_limit
Lucia Hradecká's avatar
Lucia Hradecká committed
        }


class AffineTransform(DualTransform):
Lucia D. Hradecka's avatar
Lucia D. Hradecka committed
    """Affine transformation of the input image with given parameters. Image shape remains unchanged.
Lucia Hradecká's avatar
Lucia Hradecká committed

        Args:
            angles (Tuple[float], optional): Angles of rotation for the spatial axes.

                Must be: ``(A_Z, A_Y, A_X)``.

                Defaults to ``(0, 0, 0)``.
            translation (Tuple[float], optional): Translation vector for the spatial axes.

                Must be: ``(T_Z, T_Y, T_X)``.

                Defaults to ``(0, 0, 0)``.
            scale (Tuple[float], optional): Scales for the spatial axes.

                Must be: ``(S_Z, S_Y, S_X)``.

                Defaults to ``(1, 1, 1)``.
            spacing (Tuple[float, float, float], optional): Voxel spacing for individual spatial dimensions.

                Must be: ``(S1, S2, S3)`` (a scale for each spatial dimension must be given).

                Defaults to ``(1, 1, 1)``.
            change_to_isotropic (bool, optional): Change data from anisotropic to isotropic.

                Defaults to ``False``.
            interpolation (str, optional): SimpleITK interpolation type for `image` and `float_mask`.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Must be one of ``linear``, ``nearest``, ``bspline``, ``gaussian``.

                For `mask`, the ``nearest`` interpolation is always used.

                Defaults to ``linear``.
Lucia Hradecká's avatar
Lucia Hradecká committed
            border_mode (str, optional): Values outside image domain are filled according to this mode.

                Defaults to ``'constant'``.
            ival (float, optional): Value of `image` voxels outside of the `image` domain. Only applied when ``border_mode = 'constant'``.

                Defaults to ``0``.
            mval (float, optional): Value of `mask` and `float_mask` voxels outside of the domain. Only applied when ``border_mode = 'constant'``.
Lucia Hradecká's avatar
Lucia Hradecká committed

                Defaults to ``0``.