torchio.transforms.augmentation.intensity.random_noise

Complexity

Total Complexity

16

Size/Duplication

Total Lines	142
Duplicated Lines	0 %

Importance

Changes

0

8 Methods

Rating	Name	Duplication	Size	Complexity
A	RandomNoise.__init__()	0	10	1
A	Noise.get_arguments()	0	2	1
A	Noise.apply_transform()	0	10	4
A	RandomNoise.apply_transform()	0	10	2
A	Noise.__init__()	0	13	1
A	RandomNoise.get_params()	0	9	1
A	Noise.parse_arguments()	0	12	3
A	Noise.inverse()	0	4	1

import copy
from collections import defaultdict
from typing import Tuple, Optional, Union, List, Dict, Sequence

import torch
from ....torchio import DATA
from ....utils import get_random_seed
from ....data.subject import Subject
from ... import IntensityTransform
from .. import RandomTransform


class RandomNoise(RandomTransform, IntensityTransform):
    r"""Add Gaussian noise with random parameters.

    Add noise sampled from a normal distribution with random parameters.

    Args:
        mean: Mean :math:`\mu` of the Gaussian distribution
            from which the noise is sampled.
            If two values :math:`(a, b)` are provided,
            then :math:`\mu \sim \mathcal{U}(a, b)`.
            If only one value :math:`d` is provided,
            :math:`\mu \sim \mathcal{U}(-d, d)`.
        std: Standard deviation :math:`\sigma` of the Gaussian distribution
            from which the noise is sampled.
            If two values :math:`(a, b)` are provided,
            then :math:`\sigma \sim \mathcal{U}(a, b)`.
            If only one value :math:`d` is provided,
            :math:`\sigma \sim \mathcal{U}(0, d)`.
        p: Probability that this transform will be applied.
        seed: See :py:class:`~torchio.transforms.augmentation.RandomTransform`.
        keys: See :py:class:`~torchio.transforms.Transform`.
    """
    def __init__(
            self,
            mean: Union[float, Tuple[float, float]] = 0,
            std: Union[float, Tuple[float, float]] = (0, 0.25),
            p: float = 1,
            keys: Optional[Sequence[str]] = None,
            ):
        super().__init__(p=p, keys=keys)
        self.mean_range = self.parse_range(mean, 'mean')
        self.std_range = self.parse_range(std, 'std', min_constraint=0)

    def apply_transform(self, subject: Subject) -> Subject:
        arguments = defaultdict(dict)
        for image_name in self.get_images_dict(subject):
            mean, std, seed = self.get_params(self.mean_range, self.std_range)
            arguments['mean'][image_name] = mean
            arguments['std'][image_name] = std
            arguments['seed'][image_name] = seed
        transform = Noise(**arguments)
        transformed = transform(subject)
        return transformed

    @staticmethod
    def get_params(
            mean_range: Tuple[float, float],
            std_range: Tuple[float, float],
            ) -> Tuple[float, float]:
        mean = torch.FloatTensor(1).uniform_(*mean_range).item()
        std = torch.FloatTensor(1).uniform_(*std_range).item()
        seed = get_random_seed()
        return mean, std, seed


class Noise(IntensityTransform):
    r"""Add Gaussian noise.

    Add noise sampled from a normal distribution.

    Args:
        mean: Mean :math:`\mu` of the Gaussian distribution
            from which the noise is sampled.
        std: Standard deviation :math:`\sigma` of the Gaussian distribution
            from which the noise is sampled.
        seed:
        keys: See :py:class:`~torchio.transforms.Transform`.
    """
    def __init__(
            self,
            mean: Union[float, Dict[str, float]],
            std: Union[float, Dict[str, float]],
            seed: Union[int, Sequence[int]],
            keys: Optional[List[str]] = None,
            ):
        super().__init__(keys=keys)
        self.mean = mean
        self.std = std
        self.seed = seed
        self._has_dicts = self.parse_arguments()
        self.invert_transform = False

    def get_arguments(self):
        return {'mean': self.mean, 'std': self.std, 'seed': self.seed}

    def parse_arguments(self):
        mean_dict = isinstance(self.mean, dict)
        std_dict = isinstance(self.std, dict)
        seed_dict = isinstance(self.seed, dict)
        three_bools = mean_dict, std_dict, seed_dict
        if all(three_bools):
            return True
        elif not any(three_bools):
            return False
        else:
            message = 'All arguments must have the same type: float or dict'
            raise ValueError(message)

    def apply_transform(self, subject: Subject) -> Subject:
        args = self.mean, self.std, self.seed
        for name, image in self.get_images_dict(subject).items():
            if self._has_dicts:
                args = self.mean[name], self.std[name], self.seed[name]
            noise = get_noise(image[DATA], *args)

The variable DATA does not seem to be defined.

            if self.invert_transform:
                noise *= -1
            image[DATA] = image[DATA] + noise
        return subject

    def inverse(self):
        new = copy.deepcopy(self)
        new.invert_transform = not self.invert_transform
        return new


def get_randn(shape: Sequence[int], seed: int) -> torch.Tensor:
    torch_rng_state = torch.random.get_rng_state()
    torch.manual_seed(seed)
    noise = torch.randn(*shape)
    torch.random.set_rng_state(torch_rng_state)
    return noise


def get_noise(
        tensor: torch.Tensor,
        mean: float,
        std: float,
        seed: Optional[int] = None) -> torch.Tensor:
    return get_randn(tensor.shape, seed=seed) * std + mean