Passed
Pull Request — master (#207)
by Fernando
01:35
created

  A

Complexity

Conditions 4

Size

Total Lines 12
Code Lines 10

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 4
eloc 10
nop 3
dl 0
loc 12
rs 9.9
c 0
b 0
f 0
1
import warnings
2
from typing import Tuple, Optional, Union
3
import torch
4
import numpy as np
5
import SimpleITK as sitk
6
from ....torchio import DATA, AFFINE
7
from ....data.subject import Subject
8
from .. import RandomTransform
9
10
11
class RandomGhosting(RandomTransform):
12
    r"""Add random MRI ghosting artifact.
13
14
    Args:
15
        num_ghosts: Number of 'ghosts' :math:`n` in the image.
16
            If :py:attr:`num_ghosts` is a tuple :math:`(a, b)`, then
17
            :math:`n \sim \mathcal{U}(a, b) \cap \mathbb{N}`.
18
        axes: Axis along which the ghosts will be created. If
19
            :py:attr:`axes` is a tuple, the axis will be randomly chosen
20
            from the passed values.
21
        intensity: Positive number representing the artifact strength
22
            :math:`s` with respect to the maximum of the :math:`k`-space.
23
            If ``0``, the ghosts will not be visible. If a tuple
24
            :math:`(a, b)`, is provided then
25
            :math:`s \sim \mathcal{U}(a, b)`.
26
        restore: Number between ``0`` and ``1`` indicating how much of the
27
            :math:`k`-space center should be restored after removing the planes
28
            that generate the artifact.
29
        p: Probability that this transform will be applied.
30
        seed: See :py:class:`~torchio.transforms.augmentation.RandomTransform`.
31
32
    .. note:: The execution time of this transform does not depend on the
33
        number of ghosts.
34
    """
35
    def __init__(
36
            self,
37
            num_ghosts: Union[int, Tuple[int, int]] = (4, 10),
38
            axes: Union[int, Tuple[int, ...]] = (0, 1, 2),
39
            intensity: Union[float, Tuple[float, float]] = (0.5, 1),
40
            restore: float = 0.02,
41
            p: float = 1,
42
            seed: Optional[int] = None,
43
            ):
44
        super().__init__(p=p, seed=seed)
45
        if not isinstance(axes, tuple):
46
            try:
47
                axes = tuple(axes)
48
            except TypeError:
49
                axes = (axes,)
50
        for axis in axes:
51
            if axis not in (0, 1, 2):
52
                raise ValueError(f'Axes must be in (0, 1, 2), not "{axes}"')
53
        self.axes = axes
54
        if isinstance(num_ghosts, int):
55
            self.num_ghosts_range = num_ghosts, num_ghosts
56
        elif isinstance(num_ghosts, tuple) and len(num_ghosts) == 2:
57
            self.num_ghosts_range = num_ghosts
58
        self.intensity_range = self.parse_range(intensity, 'intensity')
59
        for n in self.intensity_range:
60
            if n < 0:
61
                message = (
62
                    f'Intensity must be a positive number, not {n}')
63
                raise ValueError(message)
64
        if not 0 <= restore < 1:
65
            message = (
66
                f'Restore must be a number between 0 and 1, not {restore}')
67
            raise ValueError(message)
68
        self.restore = restore
69
70
    def apply_transform(self, sample: Subject) -> dict:
71
        random_parameters_images_dict = {}
72
        for image_name, image_dict in sample.get_images_dict().items():
73
            data = image_dict[DATA]
0 ignored issues
show
Comprehensibility Best Practice introduced by
The variable DATA does not seem to be defined.
Loading history...
74
            is_2d = data.shape[-3] == 1
75
            axes = [a for a in self.axes if a != 0] if is_2d else self.axes
76
            params = self.get_params(
77
                self.num_ghosts_range,
78
                axes,
79
                self.intensity_range,
80
            )
81
            num_ghosts_param, axis_param, intensity_param = params
82
            random_parameters_dict = {
83
                'axis': axis_param,
84
                'num_ghosts': num_ghosts_param,
85
                'intensity': intensity_param,
86
            }
87
            random_parameters_images_dict[image_name] = random_parameters_dict
88
            image_dict[DATA][0] = self.add_artifact(
89
                data[0],
90
                num_ghosts_param,
91
                axis_param,
92
                intensity_param,
93
                self.restore,
94
            )
95
        sample.add_transform(self, random_parameters_images_dict)
96
        return sample
97
98
    @staticmethod
99
    def get_params(
100
            num_ghosts_range: Tuple[int, int],
101
            axes: Tuple[int, ...],
102
            intensity_range: Tuple[float, float],
103
            ) -> Tuple:
104
        ng_min, ng_max = num_ghosts_range
105
        num_ghosts = torch.randint(ng_min, ng_max + 1, (1,)).item()
106
        axis = axes[torch.randint(0, len(axes), (1,))]
107
        intensity = torch.FloatTensor(1).uniform_(*intensity_range).item()
108
        return num_ghosts, axis, intensity
109
110
    def add_artifact(
111
            self,
112
            tensor: torch.Tensor,
113
            num_ghosts: int,
114
            axis: int,
115
            intensity: float,
116
            restore_center: float,
117
            ):
118
        array = tensor.numpy()
119
        spectrum = self.fourier_transform(array)
120
121
        ri, rj, rk = np.round(restore_center * np.array(array.shape)).astype(np.uint16)
122
        mi, mj, mk = np.array(array.shape) // 2
123
124
        # Variable "planes" is the part the spectrum that will be modified
125
        if axis == 0:
126
            planes = spectrum[::num_ghosts, :, :]
127
            restore = spectrum[mi, :, :].copy()
128
        elif axis == 1:
129
            planes = spectrum[:, ::num_ghosts, :]
130
            restore = spectrum[:, mj, :].copy()
131
        elif axis == 2:
132
            planes = spectrum[:, :, ::num_ghosts]
133
            restore = spectrum[:, :, mk].copy()
134
135
        # Multiply by 0 if intensity is 1
136
        planes *= 1 - intensity
0 ignored issues
show
introduced by
The variable planes does not seem to be defined for all execution paths.
Loading history...
137
138
        # Restore the center of k-space to avoid extreme artifacts
139
        if axis == 0:
140
            spectrum[mi, :, :] = restore
0 ignored issues
show
introduced by
The variable restore does not seem to be defined for all execution paths.
Loading history...
141
        elif axis == 1:
142
            spectrum[:, mj, :] = restore
143
        elif axis == 2:
144
            spectrum[:, :, mk] = restore
145
146
        array_ghosts = self.inv_fourier_transform(spectrum)
147
        array_ghosts = np.real(array_ghosts)
148
        return torch.from_numpy(array_ghosts)
149