Passed
Pull Request — master (#246)
by Fernando
01:04
created

torchio.transforms.augmentation.intensity.random_ghosting   A

Complexity

Total Complexity 22

Size/Duplication

Total Lines 165
Duplicated Lines 0 %

Importance

Changes 0
Metric Value
wmc 22
eloc 112
dl 0
loc 165
rs 10
c 0
b 0
f 0

5 Methods

Rating   Name   Duplication   Size   Complexity  
C RandomGhosting.add_artifact() 0 43 9
A RandomGhosting.apply_transform() 0 31 4
A RandomGhosting.parse_restore() 0 9 3
B RandomGhosting.__init__() 0 24 5
A RandomGhosting.get_params() 0 11 1
1
from typing import Tuple, Optional, Union
2
import torch
3
import numpy as np
4
from ....torchio import DATA
5
from ....data.subject import Subject
6
from .. import RandomTransform
7
8
9
class RandomGhosting(RandomTransform):
10
    r"""Add random MRI ghosting artifact.
11
12
    Discrete "ghost" artifacts may occur along the phase-encode direction
13
    whenever the position or signal intensity of imaged structures within the
14
    field-of-view vary or move in a regular (periodic) fashion. Pulsatile flow
15
    of blood or CSF, cardiac motion, and respiratory motion are the most
16
    important patient-related causes of ghost artifacts in clinical MR imaging
17
    (from `mriquestions.com <http://mriquestions.com/why-discrete-ghosts.html>`_).
18
19
    Args:
20
        num_ghosts: Number of 'ghosts' :math:`n` in the image.
21
            If :py:attr:`num_ghosts` is a tuple :math:`(a, b)`, then
22
            :math:`n \sim \mathcal{U}(a, b) \cap \mathbb{N}`.
23
            If only one value :math:`d` is provided,
24
            :math:`n \sim \mathcal{U}(0, d) \cap \mathbb{N}`.
25
        axes: Axis along which the ghosts will be created. If
26
            :py:attr:`axes` is a tuple, the axis will be randomly chosen
27
            from the passed values.
28
        intensity: Positive number representing the artifact strength
29
            :math:`s` with respect to the maximum of the :math:`k`-space.
30
            If ``0``, the ghosts will not be visible. If a tuple
31
            :math:`(a, b)` is provided then :math:`s \sim \mathcal{U}(a, b)`.
32
            If only one value :math:`d` is provided,
33
            :math:`s \sim \mathcal{U}(0, d)`.
34
        restore: Number between ``0`` and ``1`` indicating how much of the
35
            :math:`k`-space center should be restored after removing the planes
36
            that generate the artifact.
37
        p: Probability that this transform will be applied.
38
        seed: See :py:class:`~torchio.transforms.augmentation.RandomTransform`.
39
40
    .. note:: The execution time of this transform does not depend on the
41
        number of ghosts.
42
    """
43
    def __init__(
44
            self,
45
            num_ghosts: Union[int, Tuple[int, int]] = (4, 10),
46
            axes: Union[int, Tuple[int, ...]] = (0, 1, 2),
47
            intensity: Union[float, Tuple[float, float]] = (0.5, 1),
48
            restore: float = 0.02,
49
            p: float = 1,
50
            seed: Optional[int] = None,
51
            ):
52
        super().__init__(p=p, seed=seed)
53
        if not isinstance(axes, tuple):
54
            try:
55
                axes = tuple(axes)
56
            except TypeError:
57
                axes = (axes,)
58
        for axis in axes:
59
            if axis not in (0, 1, 2):
60
                raise ValueError(f'Axes must be in (0, 1, 2), not "{axes}"')
61
        self.axes = axes
62
        self.num_ghosts_range = self.parse_range(
63
            num_ghosts, 'num_ghosts', min_constraint=0, type_constraint=int)
64
        self.intensity_range = self.parse_range(
65
            intensity, 'intensity_range', min_constraint=0)
66
        self.restore = self.parse_restore(restore)
67
68
    @staticmethod
69
    def parse_restore(restore):
70
        if not isinstance(restore, float):
71
            raise TypeError(f'Restore must be a float, not {restore}')
72
        if not 0 <= restore <= 1:
73
            message = (
74
                f'Restore must be a number between 0 and 1, not {restore}')
75
            raise ValueError(message)
76
        return restore
77
78
    def apply_transform(self, sample: Subject) -> dict:
79
        random_parameters_images_dict = {}
80
        for image_name, image in sample.get_images_dict().items():
81
            transformed_tensors = []
82
            is_2d = image.is_2d()
83
            axes = [a for a in self.axes if a != 0] if is_2d else self.axes
84
            for channel_idx, tensor in enumerate(image[DATA]):
0 ignored issues
show
Comprehensibility Best Practice introduced by
The variable DATA does not seem to be defined.
Loading history...
85
                params = self.get_params(
86
                    self.num_ghosts_range,
87
                    axes,
88
                    self.intensity_range,
89
                )
90
                num_ghosts_param, axis_param, intensity_param = params
91
                random_parameters_dict = {
92
                    'axis': axis_param,
93
                    'num_ghosts': num_ghosts_param,
94
                    'intensity': intensity_param,
95
                }
96
                key = f'{image_name}_channel_{channel_idx}'
97
                random_parameters_images_dict[key] = random_parameters_dict
98
                transformed_tensor = self.add_artifact(
99
                    tensor,
100
                    num_ghosts_param,
101
                    axis_param,
102
                    intensity_param,
103
                    self.restore,
104
                )
105
                transformed_tensors.append(transformed_tensor)
106
            image[DATA] = torch.stack(transformed_tensors)
107
        sample.add_transform(self, random_parameters_images_dict)
108
        return sample
109
110
    @staticmethod
111
    def get_params(
112
            num_ghosts_range: Tuple[int, int],
113
            axes: Tuple[int, ...],
114
            intensity_range: Tuple[float, float],
115
            ) -> Tuple:
116
        ng_min, ng_max = num_ghosts_range
117
        num_ghosts = torch.randint(ng_min, ng_max + 1, (1,)).item()
118
        axis = axes[torch.randint(0, len(axes), (1,))]
119
        intensity = torch.FloatTensor(1).uniform_(*intensity_range).item()
120
        return num_ghosts, axis, intensity
121
122
    def add_artifact(
123
            self,
124
            tensor: torch.Tensor,
125
            num_ghosts: int,
126
            axis: int,
127
            intensity: float,
128
            restore_center: float,
129
            ):
130
        if not num_ghosts or not intensity:
131
            return tensor
132
133
        array = tensor.numpy()
134
        spectrum = self.fourier_transform(array)
135
136
        shape = np.array(array.shape)
137
        ri, rj, rk = np.round(restore_center * shape).astype(np.uint16)
138
        mi, mj, mk = np.array(array.shape) // 2
139
140
        # Variable "planes" is the part of the spectrum that will be modified
141
        if axis == 0:
142
            planes = spectrum[::num_ghosts, :, :]
143
            restore = spectrum[mi, :, :].copy()
144
        elif axis == 1:
145
            planes = spectrum[:, ::num_ghosts, :]
146
            restore = spectrum[:, mj, :].copy()
147
        elif axis == 2:
148
            planes = spectrum[:, :, ::num_ghosts]
149
            restore = spectrum[:, :, mk].copy()
150
151
        # Multiply by 0 if intensity is 1
152
        planes *= 1 - intensity
0 ignored issues
show
introduced by
The variable planes does not seem to be defined for all execution paths.
Loading history...
153
154
        # Restore the center of k-space to avoid extreme artifacts
155
        if axis == 0:
156
            spectrum[mi, :, :] = restore
0 ignored issues
show
introduced by
The variable restore does not seem to be defined for all execution paths.
Loading history...
157
        elif axis == 1:
158
            spectrum[:, mj, :] = restore
159
        elif axis == 2:
160
            spectrum[:, :, mk] = restore
161
162
        array_ghosts = self.inv_fourier_transform(spectrum)
163
        array_ghosts = np.real(array_ghosts)
164
        return torch.from_numpy(array_ghosts)
165