torchio.transforms.augmentation.intensity.random_labels_to_image

Complexity

Total Complexity

42

Size/Duplication

Total Lines	263
Duplicated Lines	0 %

Importance

Changes

0

9 Methods

Rating	Name	Duplication	Size	Complexity
A	RandomLabelsToImage.parse_default_gaussian_parameters()	0	12	3
A	RandomLabelsToImage.get_params()	0	12	2
A	RandomLabelsToImage.__init__()	0	17	1
A	RandomLabelsToImage.parse_gaussian_parameter()	0	18	4
B	RandomLabelsToImage.parse_gaussian_parameters()	0	23	7
C	RandomLabelsToImage.apply_transform()	0	49	9
A	RandomLabelsToImage.parse_pv_label_maps()	0	14	4
A	RandomLabelsToImage.generate_tissue()	0	10	1
C	RandomLabelsToImage.parse_keys()	0	19	11

from typing import Union, Tuple, Optional, Dict, Sequence
import torch
import numpy as np
from ....torchio import DATA, TypeData, TypeRangeFloat, TypeNumber, AFFINE, INTENSITY
from ....data.subject import Subject
from ....data.image import ScalarImage
from .. import RandomTransform




class RandomLabelsToImage(RandomTransform):
    MEAN_RANGE = (0.1, 0.9)
    STD_RANGE = (0.01, 0.1)
    r"""Generate an image from a segmentation.

    Based on the work by `Billot et al., A Learning Strategy for
    Contrast-agnostic MRI Segmentation <https://arxiv.org/abs/2003.01995>`_.

    Args:
        label_key: String designating the label map in the sample
            that will be used to generate the new image.
            Cannot be set at the same time as :py:attr:`pv_label_keys`.
        pv_label_keys: Sequence of strings designating the partial-volume (PV)
            label maps in the sample that will be used to generate the new
            image. Cannot be set at the same time as :py:attr:`label_key`.
        image_key: String designating the key to which the new volume will be
            saved. If this key corresponds to an already existing volume,
            voxels that have a value of 0 in the label maps will be filled with
            the corresponding values in the original volume.
        gaussian_parameters: Dictionary containing the mean and standard
            deviation for each label. For each value :math:`v`, if a tuple
            :math:`(a, b)` is provided then :math:`v \sim \mathcal{U}(a, b)`.
            If no value is given for a label, the value from
            :py:attr:`default_gaussian_parameters` will be used.
        default_gaussian_parameters: Dictionary containing the default
            mean and standard deviation used for all labels that are not
            defined in :py:attr:`gaussian_parameters`.
            Default values are ``(0.1, 0.9)`` for the mean and
            ``(0.01, 0.1)`` for the standard deviation.
        binarize: If ``True``, PV label maps will be binarized.
            Does not have any effects if not using PV label maps.
            Binarization is done taking the highest value per voxel
            in the different PV label maps.
        p: Probability that this transform will be applied.
        seed: See :py:class:`~torchio.transforms.augmentation.RandomTransform`.

    .. note:: It is recommended to blur the new images to make the result more
        realistic. See
        :py:class:`~torchio.transforms.augmentation.intensity.random_blur.RandomBlur`.

    Example:
        >>> import torchio
        >>> from torchio import RandomLabelsToImage, DATA, RescaleIntensity, Compose
        >>> from torchio.datasets import Colin27
        >>> colin = Colin27(2008)
        >>> # Using the default gaussian_parameters
        >>> transform = RandomLabelsToImage(label_key='cls')
        >>> # Using custom gaussian_parameters
        >>> label_values = colin['cls'][DATA].unique().round().long()
        >>> gaussian_parameters = {
        ...     label: {
        ...         'mean': i / len(label_values),
        ...         'std': 0.01
        ...     }
        ...     for i, label in enumerate(label_values)
        ... }
        >>> transform = RandomLabelsToImage(label_key='cls', gaussian_parameters=gaussian_parameters)
        >>> transformed = transform(colin)  # colin has a new key 'image' with the simulated image
        >>> # Filling holes of the simulated image with the original T1 image
        >>> rescale_transform = RescaleIntensity((0, 1), (1, 99))   # Rescale intensity before filling holes
        >>> simulation_transform = RandomLabelsToImage(
        ...     label_key='cls',
        ...     image_key='t1',
        ...     gaussian_parameters={0: {'mean': 0, 'std': 0}}
        ... )
        >>> transform = Compose([rescale_transform, simulation_transform])
        >>> transformed = transform(colin)  # colin's key 't1' has been replaced with the simulated image
    """
    def __init__(
            self,
            label_key: Optional[str] = None,
            pv_label_keys: Optional[Sequence[str]] = None,
            image_key: str = 'image',
            gaussian_parameters: Optional[Dict[Union[str, TypeNumber], Dict[str, TypeRangeFloat]]] = None,
            default_gaussian_parameters: Optional[Dict[str, TypeRangeFloat]] = None,
            binarize: bool = False,
            p: float = 1,
            seed: Optional[int] = None,
            ):
        super().__init__(p=p, seed=seed)
        self.label_key, self.pv_label_keys = self.parse_keys(label_key, pv_label_keys)
        self.default_gaussian_parameters = self.parse_default_gaussian_parameters(default_gaussian_parameters)
        self.gaussian_parameters = self.parse_gaussian_parameters(gaussian_parameters)
        self.image_key = image_key
        self.binarize = binarize

    @staticmethod
    def parse_keys(label_key, pv_label_keys):
        if label_key is not None and pv_label_keys is not None:
            raise ValueError('"label_key" and "pv_label_keys" can\'t be set at the same time.')
        if label_key is None and pv_label_keys is None:
            raise ValueError('One of "label_key" and "pv_label_keys" must be set.')
        if label_key is not None and not isinstance(label_key, str):
            raise TypeError(f'"label_key" must be a string, not {label_key}')
        if pv_label_keys is not None:
            try:
                iter(pv_label_keys)
            except TypeError:
                raise TypeError(f'"pv_label_keys" must be a sequence of strings, not {pv_label_keys}')
            for key in pv_label_keys:
                if not isinstance(key, str):
                    raise TypeError(f'Every key of "pv_label_keys" must be a string, found {key}')
            pv_label_keys = list(pv_label_keys)

        return label_key, pv_label_keys

    def apply_transform(self, sample: Subject) -> dict:
        random_parameters_images_dict = {}
        original_image = sample.get(self.image_key)

        if self.pv_label_keys is not None:
            label_map, affine = self.parse_pv_label_maps(self.pv_label_keys, sample)
            n_labels, *image_shape = label_map.shape
            labels = self.pv_label_keys
            values = list(range(n_labels))

            if self.binarize:
                # Take label with highest value in voxel
                max_label, label_map = label_map.max(dim=0)
                # Remove values where all labels are 0
                label_map[max_label == 0] = -1

        else:
            label_map = sample[self.label_key][DATA][0]

The variable DATA does not seem to be defined.

            affine = sample[self.label_key][AFFINE]
            image_shape = label_map.shape
            values = label_map.unique()
            labels = [int(key) for key in values.round()]

        tissues = torch.zeros(image_shape)

        for i, label in enumerate(labels):
            mean, std = self.get_params(label)
            if self.pv_label_keys is not None and not self.binarize:
                mask = label_map[i]
            else:
                mask = label_map == values[i]
            tissues += self.generate_tissue(mask, mean, std)

            random_parameters_images_dict[label] = {
                'mean': mean,
                'std': std
            }

        final_image = ScalarImage(affine=affine, tensor=tissues)

        if original_image is not None:
            if self.pv_label_keys is not None and not self.binarize:
                label_map = label_map.sum(dim=0)
            bg_mask = label_map.unsqueeze(0) <= 0
            final_image[DATA][bg_mask] = original_image[DATA][bg_mask]

        sample.add_image(final_image, self.image_key)
        sample.add_transform(self, random_parameters_images_dict)
        return sample

    def parse_default_gaussian_parameters(self, default_gaussian_parameters):
        if default_gaussian_parameters is None:
            return {'mean': self.MEAN_RANGE, 'std': self.STD_RANGE}

        if list(default_gaussian_parameters.keys()) != ['mean', 'std']:
            raise KeyError(f'Default gaussian parameters {default_gaussian_parameters.keys()} do not '
                           f'match {["mean", "std"]}')

        mean = self.parse_gaussian_parameter(default_gaussian_parameters['mean'], 'mean')
        std = self.parse_gaussian_parameter(default_gaussian_parameters['std'], 'std')

        return {'mean': mean, 'std': std}

    def parse_gaussian_parameters(self, gaussian_parameters):
        if gaussian_parameters is None:
            gaussian_parameters = {}

        if self.pv_label_keys is not None:
            if not set(self.pv_label_keys).issuperset(gaussian_parameters.keys()):
                raise KeyError(f'Found keys in gaussian parameters {gaussian_parameters.keys()} '
                               f'not in pv_label_keys {self.pv_label_keys}')

        parsed_gaussian_parameters = {}

        for label_key, dictionary in gaussian_parameters.items():
            if 'mean' in dictionary:
                mean = self.parse_gaussian_parameter(dictionary['mean'], 'mean')
            else:
                mean = self.default_gaussian_parameters['mean']
            if 'std' in dictionary:
                std = self.parse_gaussian_parameter(dictionary['std'], 'std')
            else:
                std = self.default_gaussian_parameters['std']
            parsed_gaussian_parameters.update({label_key: {'mean': mean, 'std': std}})

        return parsed_gaussian_parameters

    @staticmethod
    def parse_gaussian_parameter(
            nums_range: TypeRangeFloat,
            name: str,
            ) -> Tuple[float, float]:
        if isinstance(nums_range, (int, float)):
            return nums_range, nums_range

        if len(nums_range) != 2:
            raise ValueError(
                f'If {name} is a sequence,'
                f' it must be of len 2, not {nums_range}')
        min_value, max_value = nums_range
        if min_value > max_value:
            raise ValueError(
                f'If {name} is a sequence, the second value must be'
                f' equal or greater than the first, not {nums_range}')
        return min_value, max_value

    @staticmethod
    def parse_pv_label_maps(
            pv_label_keys: Sequence[str],
            sample: dict,
            ) -> (TypeData, TypeData):
        try:
            label_map = torch.cat([sample[key][DATA] for key in pv_label_keys], dim=0)
        except RuntimeError:
            raise RuntimeError('PV label maps have different shapes, make sure they all have the same shapes.')
        affine = sample[pv_label_keys[0]][AFFINE]
        for key in pv_label_keys[1:]:
            if not np.array_equal(affine, sample[key][AFFINE]):
                raise RuntimeWarning('Be careful, PV label maps with different affines were found.')
        return label_map, affine

    def get_params(
            self,
            label: Union[str, TypeNumber]
            ) -> Tuple[float, float]:
        if label in self.gaussian_parameters:
            mean_range, std_range = self.gaussian_parameters[label]['mean'], self.gaussian_parameters[label]['std']
        else:
            mean_range, std_range = self.default_gaussian_parameters['mean'], self.default_gaussian_parameters['std']

        mean = torch.FloatTensor(1).uniform_(*mean_range).item()
        std = torch.FloatTensor(1).uniform_(*std_range).item()
        return mean, std

    @staticmethod
    def generate_tissue(
            data: TypeData,
            mean: TypeNumber,
            std: TypeNumber,
            ) -> TypeData:
        # Create the simulated tissue using a gaussian random variable
        data_shape = data.shape
        gaussian = torch.randn(data_shape) * std + mean
        return gaussian * data