|
1
|
|
|
from numbers import Number |
|
2
|
|
|
from typing import Callable, Union, List, Optional |
|
3
|
|
|
|
|
4
|
|
|
import numpy as np |
|
5
|
|
|
import nibabel as nib |
|
6
|
|
|
import torch |
|
7
|
|
|
|
|
8
|
|
|
from ....torchio import DATA, AFFINE |
|
9
|
|
|
from ....data.subject import Subject |
|
10
|
|
|
from .bounds_transform import BoundsTransform, TypeBounds |
|
11
|
|
|
|
|
12
|
|
|
|
|
13
|
|
|
class Pad(BoundsTransform): |
|
14
|
|
|
r"""Pad an image. |
|
15
|
|
|
|
|
16
|
|
|
Args: |
|
17
|
|
|
padding: Tuple |
|
18
|
|
|
:math:`(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})` |
|
19
|
|
|
defining the number of values padded to the edges of each axis. |
|
20
|
|
|
If the initial shape of the image is |
|
21
|
|
|
:math:`W \times H \times D`, the final shape will be |
|
22
|
|
|
:math:`(w_{ini} + W + w_{fin}) \times (h_{ini} + H + h_{fin}) |
|
23
|
|
|
\times (d_{ini} + D + d_{fin})`. |
|
24
|
|
|
If only three values :math:`(w, h, d)` are provided, then |
|
25
|
|
|
:math:`w_{ini} = w_{fin} = w`, |
|
26
|
|
|
:math:`h_{ini} = h_{fin} = h` and |
|
27
|
|
|
:math:`d_{ini} = d_{fin} = d`. |
|
28
|
|
|
If only one value :math:`n` is provided, then |
|
29
|
|
|
:math:`w_{ini} = w_{fin} = h_{ini} = h_{fin} = |
|
30
|
|
|
d_{ini} = d_{fin} = n`. |
|
31
|
|
|
padding_mode: See possible modes in `NumPy docs`_. If it is a number, |
|
32
|
|
|
the mode will be set to ``'constant'``. |
|
33
|
|
|
p: Probability that this transform will be applied. |
|
34
|
|
|
keys: See :py:class:`~torchio.transforms.Transform`. |
|
35
|
|
|
|
|
36
|
|
|
.. _NumPy docs: https://numpy.org/doc/stable/reference/generated/numpy.pad.html |
|
37
|
|
|
""" |
|
38
|
|
|
|
|
39
|
|
|
PADDING_MODES = ( |
|
40
|
|
|
'empty', |
|
41
|
|
|
'edge', |
|
42
|
|
|
'wrap', |
|
43
|
|
|
'constant', |
|
44
|
|
|
'linear_ramp', |
|
45
|
|
|
'maximum', |
|
46
|
|
|
'mean', |
|
47
|
|
|
'median', |
|
48
|
|
|
'minimum', |
|
49
|
|
|
'reflect', |
|
50
|
|
|
'symmetric', |
|
51
|
|
|
) |
|
52
|
|
|
|
|
53
|
|
|
def __init__( |
|
54
|
|
|
self, |
|
55
|
|
|
padding: TypeBounds, |
|
56
|
|
|
padding_mode: Union[str, float] = 0, |
|
57
|
|
|
p: float = 1, |
|
58
|
|
|
keys: Optional[List[str]] = None, |
|
59
|
|
|
): |
|
60
|
|
|
super().__init__(padding, p=p, keys=keys) |
|
61
|
|
|
self.padding_mode, self.fill = self.parse_padding_mode(padding_mode) |
|
62
|
|
|
|
|
63
|
|
|
@classmethod |
|
64
|
|
|
def parse_padding_mode(cls, padding_mode): |
|
65
|
|
|
if padding_mode in cls.PADDING_MODES: |
|
66
|
|
|
fill = None |
|
67
|
|
|
elif isinstance(padding_mode, Number): |
|
68
|
|
|
fill = padding_mode |
|
69
|
|
|
padding_mode = 'constant' |
|
70
|
|
|
else: |
|
71
|
|
|
message = ( |
|
72
|
|
|
f'Padding mode "{padding_mode}" not valid. Valid options are' |
|
73
|
|
|
f' {list(cls.PADDING_MODES)} or a number' |
|
74
|
|
|
) |
|
75
|
|
|
raise KeyError(message) |
|
76
|
|
|
return padding_mode, fill |
|
77
|
|
|
|
|
78
|
|
|
def apply_transform(self, subject: Subject) -> Subject: |
|
79
|
|
|
low = self.bounds_parameters[::2] |
|
80
|
|
|
for image in self.get_images(subject): |
|
81
|
|
|
new_origin = nib.affines.apply_affine(image.affine, -np.array(low)) |
|
82
|
|
|
new_affine = image.affine.copy() |
|
83
|
|
|
new_affine[:3, 3] = new_origin |
|
84
|
|
|
kwargs = dict(mode=self.padding_mode) |
|
85
|
|
|
if self.padding_mode == 'constant': |
|
86
|
|
|
kwargs['constant_values'] = self.fill |
|
87
|
|
|
pad_params = self.bounds_parameters |
|
88
|
|
|
paddings = (0, 0), pad_params[:2], pad_params[2:4], pad_params[4:] |
|
89
|
|
|
padded = np.pad(image[DATA], paddings, **kwargs) |
|
|
|
|
|
|
90
|
|
|
image[DATA] = torch.from_numpy(padded) |
|
91
|
|
|
image[AFFINE] = new_affine |
|
92
|
|
|
return subject |
|
93
|
|
|
|
fepegar /
torchio
Loading history...