forked from mlcommons/GaNDLF
-
Notifications
You must be signed in to change notification settings - Fork 0
/
wrap_torchio.py
122 lines (94 loc) · 3.6 KB
/
wrap_torchio.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import numpy as np
from torchio.transforms import (
OneOf,
RandomMotion,
RandomGhosting,
RandomSpike,
RandomAffine,
RandomElasticDeformation,
RandomBiasField,
RandomSwap,
RandomNoise,
RandomAnisotropy,
RandomFlip,
RandomGamma,
)
from .blur_enhanced import RandomBlurEnhanced
from .noise_enhanced import RandomNoiseEnhanced
## define helper functions to create transforms
## todo: ability to change interpolation type from config file
## todo: ability to change the dimensionality according to the config file
# define individual functions/lambdas for augmentations to handle properties
def mri_artifact(parameters):
return OneOf(
{RandomGhosting(): 0.5, RandomSpike(): 0.5}, p=parameters["probability"]
)
def motion(parameters):
return RandomMotion(
degrees=parameters["degrees"],
translation=parameters["translation"],
num_transforms=parameters["num_transforms"],
image_interpolation=parameters["interpolation"],
p=parameters["probability"],
)
def affine(parameters):
return RandomAffine(
scales=parameters["scales"],
degrees=parameters["degrees"],
translation=parameters["translation"],
p=parameters["probability"],
)
def elastic(parameters):
# define defaults
parameters["num_control_points"] = parameters.get("num_control_points", None)
parameters["max_displacement"] = parameters.get("max_displacement", None)
parameters["locked_borders"] = parameters.get("locked_borders", 2)
assert (
"patch_size" in parameters
), "'patch_size' must be defined for elastic deformation"
if parameters["num_control_points"] is None:
# define the control points and swap axes for augmentation
parameters["num_control_points"] = []
for _, n in enumerate(parameters["patch_size"]):
parameters["num_control_points"].append(max(n, 5)) # always at least have 5
if parameters["max_displacement"] is None:
parameters["max_displacement"] = np.divide(parameters["num_control_points"], 10)
if parameters["num_control_points"][-1] == 1:
# ensure maximum displacement is never greater than patch size
parameters["max_displacement"][-1] = 0.1
parameters["max_displacement"] = parameters["max_displacement"].tolist()
return RandomElasticDeformation(
num_control_points=parameters["num_control_points"],
max_displacement=parameters["max_displacement"],
locked_borders=parameters["locked_borders"],
p=parameters["probability"],
)
def swap(parameters):
return RandomSwap(
patch_size=parameters["patch_size"],
num_iterations=100,
p=parameters["probability"],
)
def bias(parameters):
return RandomBiasField(coefficients=0.5, order=3, p=parameters["probability"])
def blur(parameters):
return RandomBlurEnhanced(std=parameters["std"], p=parameters["probability"])
def noise(parameters):
return RandomNoise(
mean=parameters["mean"], std=parameters["std"], p=parameters["probability"]
)
def noise_var(parameters):
return RandomNoiseEnhanced(
mean=parameters["mean"], std=parameters["std"], p=parameters["probability"]
)
def gamma(parameters):
return RandomGamma(p=parameters["probability"])
def flip(parameters):
return RandomFlip(axes=parameters["axis"], p=parameters["probability"])
def anisotropy(parameters):
return RandomAnisotropy(
axes=parameters["axis"],
downsampling=parameters["downsampling"],
scalars_only=True,
p=parameters["probability"],
)