def get_moreDA_augmentation(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1,
seeds_train=None,
seeds_val=None,
order_seg=1,
order_data=3,
deep_supervision_scales=None,
soft_ds=False,
classes=None,
pin_memory=True):
assert params.get(
'mirror') is None, "old version of params, use new keyword do_mirror"
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
ignore_axes = (0, )
tr_transforms.append(Convert3DTo2DTransform())
else:
ignore_axes = None
tr_transforms.append(
SpatialTransform(patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"),
sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"),
angle_x=params.get("rotation_x"),
angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"),
p_rot_per_axis=params.get("rotation_p_per_axis"),
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=order_data,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=order_seg,
random_crop=params.get("random_crop"),
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot"),
independent_scale_for_each_axis=params.get(
"independent_scale_factor_for_each_axis")))
if params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
# we need to put the color augmentations after the dummy 2d part (if applicable). Otherwise the overloaded color
# channel gets in the way
tr_transforms.append(GaussianNoiseTransform(p_per_sample=0.1))
tr_transforms.append(
GaussianBlurTransform((0.5, 1.),
different_sigma_per_channel=True,
p_per_sample=0.2,
p_per_channel=0.5))
tr_transforms.append(
BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25),
p_per_sample=0.15))
if params.get("do_additive_brightness"):
tr_transforms.append(
BrightnessTransform(
params.get("additive_brightness_mu"),
params.get("additive_brightness_sigma"),
True,
p_per_sample=params.get("additive_brightness_p_per_sample"),
p_per_channel=params.get("additive_brightness_p_per_channel")))
tr_transforms.append(ContrastAugmentationTransform(p_per_sample=0.15))
tr_transforms.append(
SimulateLowResolutionTransform(zoom_range=(0.5, 1),
per_channel=True,
p_per_channel=0.5,
order_downsample=0,
order_upsample=3,
p_per_sample=0.25,
ignore_axes=ignore_axes))
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
True,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=0.1)) # inverted gamma
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
False,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
if params.get("do_mirror") or params.get("mirror"):
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
if params.get("mask_was_used_for_normalization") is not None:
mask_was_used_for_normalization = params.get(
"mask_was_used_for_normalization")
tr_transforms.append(
MaskTransform(mask_was_used_for_normalization,
mask_idx_in_seg=0,
set_outside_to=0))
tr_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
tr_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
if params.get(
"cascade_do_cascade_augmentations") is not None and params.get(
"cascade_do_cascade_augmentations"):
if params.get("cascade_random_binary_transform_p") > 0:
tr_transforms.append(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")),
0)),
p_per_sample=params.get(
"cascade_random_binary_transform_p"),
key="data",
strel_size=params.get(
"cascade_random_binary_transform_size"),
p_per_label=params.get(
"cascade_random_binary_transform_p_per_label")))
if params.get("cascade_remove_conn_comp_p") > 0:
tr_transforms.append(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")),
0)),
key="data",
p_per_sample=params.get("cascade_remove_conn_comp_p"),
fill_with_other_class_p=params.get(
"cascade_remove_conn_comp_max_size_percent_threshold"
),
dont_do_if_covers_more_than_X_percent=params.get(
"cascade_remove_conn_comp_fill_with_other_class_p")
))
tr_transforms.append(RenameTransform('seg', 'target', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
tr_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'target',
'target', classes))
else:
tr_transforms.append(
DownsampleSegForDSTransform2(deep_supervision_scales,
0,
0,
input_key='target',
output_key='target'))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=seeds_train,
pin_memory=pin_memory)
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
val_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
val_transforms.append(RenameTransform('seg', 'target', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
val_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'target',
'target', classes))
else:
val_transforms.append(
DownsampleSegForDSTransform2(deep_supervision_scales,
0,
0,
input_key='target',
output_key='target'))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=seeds_val,
pin_memory=pin_memory)
return batchgenerator_train, batchgenerator_val
def get_default_augmentation(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1,
pin_memory=True,
seeds_train=None,
seeds_val=None):
assert params.get(
'mirror') is None, "old version of params, use new keyword do_mirror"
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert3DTo2DTransform())
tr_transforms.append(
SpatialTransform(patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"),
sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"),
angle_x=params.get("rotation_x"),
angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"),
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=3,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=1,
random_crop=params.get("random_crop"),
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot"),
independent_scale_for_each_axis=params.get(
"independent_scale_factor_for_each_axis")))
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
False,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
if params.get("do_mirror"):
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
if params.get("mask_was_used_for_normalization") is not None:
mask_was_used_for_normalization = params.get(
"mask_was_used_for_normalization")
tr_transforms.append(
MaskTransform(mask_was_used_for_normalization,
mask_idx_in_seg=0,
set_outside_to=0))
tr_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
tr_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
if params.get("cascade_do_cascade_augmentations"
) and not None and params.get(
"cascade_do_cascade_augmentations"):
tr_transforms.append(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")), 0)),
p_per_sample=params.get(
"cascade_random_binary_transform_p"),
key="data",
strel_size=params.get(
"cascade_random_binary_transform_size")))
tr_transforms.append(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")), 0)),
key="data",
p_per_sample=params.get("cascade_remove_conn_comp_p"),
fill_with_other_class_p=params.get(
"cascade_remove_conn_comp_max_size_percent_threshold"),
dont_do_if_covers_more_than_X_percent=params.get(
"cascade_remove_conn_comp_fill_with_other_class_p")))
tr_transforms.append(RenameTransform('seg', 'target', True))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
# from batchgenerators.dataloading import SingleThreadedAugmenter
# batchgenerator_train = SingleThreadedAugmenter(dataloader_train, tr_transforms)
# import IPython;IPython.embed()
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=seeds_train,
pin_memory=pin_memory)
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
val_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
val_transforms.append(RenameTransform('seg', 'target', True))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
# batchgenerator_val = SingleThreadedAugmenter(dataloader_val, val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=seeds_val,
pin_memory=pin_memory)
return batchgenerator_train, batchgenerator_val
def get_no_augmentation(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1,
seeds_train=None,
seeds_val=None,
order_seg=1,
order_data=3,
deep_supervision_scales=None,
soft_ds=False,
classes=None,
pin_memory=True):
"""
use this instead of get_default_augmentation (drop in replacement) to turn off all data augmentation
:param dataloader_train:
:param dataloader_val:
:param patch_size:
:param params:
:param border_val_seg:
:return:
"""
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
tr_transforms.append(RemoveLabelTransform(-1, 0))
tr_transforms.append(RenameTransform('seg', 'target', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
tr_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'target',
'target', classes))
else:
tr_transforms.append(
DownsampleSegForDSTransform2(deep_supervision_scales,
0,
0,
input_key='target',
output_key='target'))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=range(params.get('num_threads')),
pin_memory=True)
batchgenerator_train.restart()
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
val_transforms.append(RenameTransform('seg', 'target', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
val_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'target',
'target', classes))
else:
val_transforms.append(
DownsampleSegForDSTransform2(deep_supervision_scales,
0,
0,
input_key='target',
output_key='target'))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=range(max(params.get('num_threads') // 2, 1)),
pin_memory=True)
batchgenerator_val.restart()
return batchgenerator_train, batchgenerator_val
def get_default_augmentation(dataloader_train, dataloader_val=None, params=None,
patch_size=None, border_val_seg=-1, pin_memory=True,
seeds_train=None, seeds_val=None, regions=None):
assert params.get('mirror') is None, "old version of params, use new keyword do_mirror"
tr_transforms = []
assert params is not None, "augmentation params expect to be not None"
if params.get("selected_data_channels") is not None:
tr_transforms.append(DataChannelSelectionTransform(params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert3DTo2DTransform())
tr_transforms.append(SpatialTransform(
patch_size, patch_center_dist_from_border=None, do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"), sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"), angle_x=params.get("rotation_x"), angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"), do_scale=params.get("do_scaling"), scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"), border_cval_data=0, order_data=3, border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=1, random_crop=params.get("random_crop"), p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"), p_rot_per_sample=params.get("p_rot"),
independent_scale_for_each_axis=params.get("independent_scale_factor_for_each_axis")
))
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"), False, True, retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
if params.get("do_mirror"):
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
tr_transforms.append(RemoveLabelTransform(-1, 0))
tr_transforms.append(RenameTransform('seg', 'target', True))
if regions is not None:
tr_transforms.append(ConvertSegmentationToRegionsTransform(regions, 'target', 'target'))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(dataloader_train, tr_transforms, params.get('num_threads'),
params.get("num_cached_per_thread"), seeds=seeds_train,
pin_memory=pin_memory)
batchgenerator_train.restart()
if dataloader_val is None:
return batchgenerator_train, None
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(DataChannelSelectionTransform(params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(SegChannelSelectionTransform(params.get("selected_seg_channels")))
val_transforms.append(RenameTransform('seg', 'target', True))
if regions is not None:
val_transforms.append(ConvertSegmentationToRegionsTransform(regions, 'target', 'target'))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(dataloader_val, val_transforms, max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"), seeds=seeds_val,
pin_memory=pin_memory)
batchgenerator_val.restart()
return batchgenerator_train, batchgenerator_val
def get_arteries_augmentation(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1,
seeds_train=None,
seeds_val=None,
order_seg=1,
order_data=3,
deep_supervision_scales=None,
soft_ds=False,
classes=None,
pin_memory=True,
regions=None,
use_nondetMultiThreadedAugmenter: bool = False):
assert params.get(
'mirror') is None, "old version of params, use new keyword do_mirror"
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
# if params.get("dummy_2D") is not None and params.get("dummy_2D"):
# ignore_axes = (0,)
# tr_transforms.append(Convert3DTo2DTransform())
# else:
# ignore_axes = None
tr_transforms.append(
SpatialTransform(patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=False,
do_rotation=False,
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=order_data,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=order_seg,
random_crop=False,
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot"),
independent_scale_for_each_axis=params.get(
"independent_scale_factor_for_each_axis")))
if params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
if params.get("do_mirror") or params.get("mirror"):
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
tr_transforms.append(RemoveLabelTransform(-1, 0))
# if params.get("move_last_seg_chanel_to_data") is not None and params.get("move_last_seg_chanel_to_data"):
# tr_transforms.append(MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"), 'seg', 'data'))
# if params.get("cascade_do_cascade_augmentations") is not None and params.get(
# "cascade_do_cascade_augmentations"):
# if params.get("cascade_random_binary_transform_p") > 0:
# tr_transforms.append(ApplyRandomBinaryOperatorTransform(
# channel_idx=list(range(-len(params.get("all_segmentation_labels")), 0)),
# p_per_sample=params.get("cascade_random_binary_transform_p"),
# key="data",
# strel_size=params.get("cascade_random_binary_transform_size"),
# p_per_label=params.get("cascade_random_binary_transform_p_per_label")))
# if params.get("cascade_remove_conn_comp_p") > 0:
# tr_transforms.append(
# RemoveRandomConnectedComponentFromOneHotEncodingTransform(
# channel_idx=list(range(-len(params.get("all_segmentation_labels")), 0)),
# key="data",
# p_per_sample=params.get("cascade_remove_conn_comp_p"),
# fill_with_other_class_p=params.get("cascade_remove_conn_comp_max_size_percent_threshold"),
# dont_do_if_covers_more_than_X_percent=params.get(
# "cascade_remove_conn_comp_fill_with_other_class_p")))
tr_transforms.append(RenameTransform('seg', 'target', True))
if regions is not None:
tr_transforms.append(
ConvertSegmentationToRegionsTransform(regions, 'target', 'target'))
# if deep_supervision_scales is not None:
# if soft_ds:
# assert classes is not None
# tr_transforms.append(DownsampleSegForDSTransform3(deep_supervision_scales, 'target', 'target', classes))
# else:
# tr_transforms.append(DownsampleSegForDSTransform2(deep_supervision_scales, 0, 0, input_key='target',
# output_key='target'))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=seeds_train,
pin_memory=pin_memory)
# batchgenerator_train = SingleThreadedAugmenter(dataloader_train, tr_transforms)
# import IPython;IPython.embed()
batchgenerator_train.restart()
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# if params.get("move_last_seg_chanel_to_data") is not None and params.get("move_last_seg_chanel_to_data"):
# val_transforms.append(MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"), 'seg', 'data'))
val_transforms.append(RenameTransform('seg', 'target', True))
if regions is not None:
val_transforms.append(
ConvertSegmentationToRegionsTransform(regions, 'target', 'target'))
# if deep_supervision_scales is not None:
# if soft_ds:
# assert classes is not None
# val_transforms.append(DownsampleSegForDSTransform3(deep_supervision_scales, 'target', 'target', classes))
# else:
# val_transforms.append(DownsampleSegForDSTransform2(deep_supervision_scales, 0, 0, input_key='target',
# output_key='target'))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=seeds_val,
pin_memory=pin_memory)
# batchgenerator_val = SingleThreadedAugmenter(dataloader_val, val_transforms)
batchgenerator_val.restart()
return batchgenerator_train, batchgenerator_val
def get_moreDA_augmentation_ae(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1,
seeds_train=None,
seeds_val=None,
order_seg=1,
order_data=3,
deep_supervision_scales=None,
soft_ds=False,
classes=None,
pin_memory=True,
anisotropy=False,
extra_label_keys=None,
val_mode=False,
use_conf=False,
global_params=None):
'''
Work as Dataloader with augmentation
:return: train_loader, val_loader
for each iterator, return {'data': (B, D, H, W), 'target': (B, D, H, W)}
'''
if global_params:
for key, value in global_params.items():
params[key] = value
print(global_params)
if not val_mode:
assert params.get(
'mirror'
) is None, "old version of params, use new keyword do_mirror"
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(
params.get("selected_seg_channels")))
# anistropic setting
if anisotropy or params.get("dummy_2D"):
ignore_axes = (0, )
tr_transforms.append(
Convert3DTo2DTransform(extra_label_keys=extra_label_keys))
patch_size = patch_size[1:] # 2D patch size
print('Using dummy2d data augmentation')
params["elastic_deform_alpha"] = (0., 200.)
params["elastic_deform_sigma"] = (9., 13.)
params["rotation_x"] = (-180. / 360 * 2. * np.pi,
180. / 360 * 2. * np.pi)
params["rotation_y"] = (-0. / 360 * 2. * np.pi,
0. / 360 * 2. * np.pi)
params["rotation_z"] = (-0. / 360 * 2. * np.pi,
0. / 360 * 2. * np.pi)
else:
ignore_axes = None
# 1. Spatial Transform: rotation, scaling
tr_transforms.append(
SpatialTransform(
patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"),
sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"),
angle_x=params.get("rotation_x"),
angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"),
p_rot_per_axis=params.get("rotation_p_per_axis"),
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=order_data,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=order_seg,
random_crop=params.get("random_crop"),
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot"),
independent_scale_for_each_axis=params.get(
"independent_scale_factor_for_each_axis"),
extra_label_keys=extra_label_keys,
do_translate=params.get("do_translate"),
p_trans=params.get("p_trans"),
trans_max_shifts=params.get("trans_max_shifts"),
trans_const_channel=params.get("trans_const_channel"),
))
if anisotropy or params.get("dummy_2D"):
tr_transforms.append(
Convert2DTo3DTransform(extra_label_keys=extra_label_keys))
'''
# 2. Noise Augmentation: gaussian noise, gaussian blur
tr_transforms.append(GaussianNoiseTransform(p_per_sample=0.1))
tr_transforms.append(GaussianBlurTransform((0.5, 1.), different_sigma_per_channel=True, p_per_sample=0.2,
p_per_channel=0.5))
# 3. Color Augmentation: brightness, constrast, low resolution, gamma_transform
tr_transforms.append(BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25), p_per_sample=0.15))
if params.get("do_additive_brightness"):
tr_transforms.append(BrightnessTransform(params.get("additive_brightness_mu"),
params.get("additive_brightness_sigma"),
True, p_per_sample=params.get("additive_brightness_p_per_sample"),
p_per_channel=params.get("additive_brightness_p_per_channel")))
tr_transforms.append(ContrastAugmentationTransform(p_per_sample=0.15))
tr_transforms.append(SimulateLowResolutionTransform(zoom_range=(0.5, 1), per_channel=True,
p_per_channel=0.5,
order_downsample=0, order_upsample=3, p_per_sample=0.25,
ignore_axes=ignore_axes))
tr_transforms.append(
GammaTransform(params.get("gamma_range"), True, True, retain_stats=params.get("gamma_retain_stats"),
p_per_sample=0.1)) # inverted gamma
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"), False, True, retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
'''
# 4. Mirror Transform
if params.get("do_mirror") or params.get("mirror"):
tr_transforms.append(
MirrorTransform(params.get("mirror_axes"),
extra_label_keys=extra_label_keys))
# if params.get("mask_was_used_for_normalization") is not None:
# mask_was_used_for_normalization = params.get("mask_was_used_for_normalization")
# tr_transforms.append(MaskTransform(mask_was_used_for_normalization, mask_idx_in_seg=0, set_outside_to=0))
tr_transforms.append(
RemoveLabelTransform(-1, 0, extra_label_keys=extra_label_keys))
tr_transforms.append(RenameTransform('data', 'image', True))
tr_transforms.append(RenameTransform('seg', 'gt', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
tr_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'gt',
'gt', classes))
else:
tr_transforms.append(
DownsampleSegForDSTransform2(
deep_supervision_scales,
0,
0,
input_key='gt',
output_key='gt',
extra_label_keys=extra_label_keys))
toTensorKeys = [
'image', 'gt'
] + extra_label_keys if extra_label_keys is not None else [
'image', 'gt'
]
tr_transforms.append(NumpyToTensor(toTensorKeys, 'float'))
tr_transforms = Compose(tr_transforms)
if seeds_train is not None:
seeds_train = [seeds_train] * params.get('num_threads')
if use_conf:
num_threads = 1
num_cached_per_thread = 1
else:
num_threads, num_cached_per_thread = params.get(
'num_threads'), params.get("num_cached_per_thread")
batchgenerator_train = MultiThreadedAugmenter(dataloader_train,
tr_transforms,
num_threads,
num_cached_per_thread,
seeds=seeds_train,
pin_memory=pin_memory)
val_transforms = []
val_transforms.append(
RemoveLabelTransform(-1, 0, extra_label_keys=extra_label_keys))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(
params.get("selected_seg_channels")))
val_transforms.append(RenameTransform('data', 'image', True))
val_transforms.append(RenameTransform('seg', 'gt', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
val_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'gt',
'gt', classes))
else:
val_transforms.append(
DownsampleSegForDSTransform2(
deep_supervision_scales,
0,
0,
input_key='gt',
output_key='gt',
extra_label_keys=extra_label_keys))
val_transforms.append(NumpyToTensor(toTensorKeys, 'float'))
val_transforms = Compose(val_transforms)
# batchgenerator_val = MultiThreadedAugmenter(dataloader_val, val_transforms, max(params.get('num_threads') // 2, 1),
# params.get("num_cached_per_thread"),
# seeds=seeds_val, pin_memory=pin_memory)
if seeds_val is not None:
seeds_val = [seeds_val] * 1
# batchgenerator_val = MultiThreadedAugmenter(dataloader_val, val_transforms, 1,
# params.get("num_cached_per_thread"),
# seeds=seeds_val, pin_memory=False)
batchgenerator_val = SingleThreadedAugmenter(dataloader_val,
val_transforms)
else:
val_transforms = []
val_transforms.append(
RemoveLabelTransform(-1, 0, extra_label_keys=extra_label_keys))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(
params.get("selected_seg_channels")))
val_transforms.append(RenameTransform('data', 'image', True))
val_transforms.append(RenameTransform('seg', 'gt', True))
if deep_supervision_scales is not None:
if soft_ds:
assert classes is not None
val_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'gt',
'gt', classes))
else:
val_transforms.append(
DownsampleSegForDSTransform2(
deep_supervision_scales,
0,
0,
input_key='gt',
output_key='gt',
extra_label_keys=extra_label_keys))
toTensorKeys = [
'image', 'gt'
] + extra_label_keys if extra_label_keys is not None else [
'image', 'gt'
]
val_transforms.append(NumpyToTensor(toTensorKeys, 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = SingleThreadedAugmenter(dataloader_val,
val_transforms)
batchgenerator_train = SingleThreadedAugmenter(dataloader_train,
val_transforms)
return batchgenerator_train, batchgenerator_val
def get_no_augmentation(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1):
"""
use this instead of get_default_augmentation (drop in replacement) to turn off all data augmentation
:param dataloader_train:
:param dataloader_val:
:param patch_size:
:param params:
:param border_val_seg:
:return:
"""
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
tr_transforms.append(RemoveLabelTransform(-1, 0))
tr_transforms.append(RenameTransform('seg', 'target', True))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=range(params.get('num_threads')),
pin_memory=True)
batchgenerator_train.restart()
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
val_transforms.append(RenameTransform('seg', 'target', True))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=range(max(params.get('num_threads') // 2, 1)),
pin_memory=True)
batchgenerator_val.restart()
return batchgenerator_train, batchgenerator_val
def get_insaneDA_augmentation(dataloader_train,
dataloader_val,
patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1,
seeds_train=None,
seeds_val=None,
order_seg=1,
order_data=3,
deep_supervision_scales=None,
soft_ds=False,
classes=None,
pin_memory=True,
regions=None):
assert params.get(
'mirror') is None, "old version of params, use new keyword do_mirror"
tr_transforms = []
# 'patch_size': array([288, 320]),
# 'border_val_seg': -1,
# 'seeds_train': None,
# 'seeds_val': None,
# 'order_seg': 1,
# 'order_data': 3,
# 'deep_supervision_scales': [[1, 1, 1],
# [1.0, 0.5, 0.5],
# [1.0, 0.25, 0.25],
# [0.5, 0.125, 0.125],
# [0.5, 0.0625, 0.0625]],
# 'soft_ds': False,
# 'classes': None,
# 'pin_memory': True,
# 'regions': None
# params
# {'selected_data_channels': None,
# 'selected_seg_channels': [0],
# 'do_elastic': True,
# 'elastic_deform_alpha': (0.0, 300.0),
# 'elastic_deform_sigma': (9.0, 15.0),
# 'p_eldef': 0.1,
# 'do_scaling': True,
# 'scale_range': (0.65, 1.6),
# 'independent_scale_factor_for_each_axis': True,
# 'p_independent_scale_per_axis': 0.3,
# 'p_scale': 0.3,
# 'do_rotation': True,
# 'rotation_x': (-3.141592653589793, 3.141592653589793),
# 'rotation_y': (-0.5235987755982988, 0.5235987755982988),
# 'rotation_z': (-0.5235987755982988, 0.5235987755982988),
# 'rotation_p_per_axis': 1,
# 'p_rot': 0.7,
# 'random_crop': False,
# 'random_crop_dist_to_border': None,
# 'do_gamma': True,
# 'gamma_retain_stats': True,
# 'gamma_range': (0.5, 1.6),
# 'p_gamma': 0.3,
# 'do_mirror': True,
# 'mirror_axes': (0, 1, 2),
# 'dummy_2D': True,
# 'mask_was_used_for_normalization': OrderedDict([(0, False)]),
# 'border_mode_data': 'constant',
# 'all_segmentation_labels': None,
# 'move_last_seg_chanel_to_data': False,
# 'cascade_do_cascade_augmentations': False,
# 'cascade_random_binary_transform_p': 0.4,
# 'cascade_random_binary_transform_p_per_label': 1,
# 'cascade_random_binary_transform_size': (1, 8),
# 'cascade_remove_conn_comp_p': 0.2,
# 'cascade_remove_conn_comp_max_size_percent_threshold': 0.15,
# 'cascade_remove_conn_comp_fill_with_other_class_p': 0.0,
# 'do_additive_brightness': True,
# 'additive_brightness_p_per_sample': 0.3,
# 'additive_brightness_p_per_channel': 1,
# 'additive_brightness_mu': 0,
# 'additive_brightness_sigma': 0.2,
# 'num_threads': 12,
# 'num_cached_per_thread': 1,
# 'patch_size_for_spatialtransform': array([288, 320])}
# selected_data_channels is None
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
# selected_seg_channels is [0]
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
# dummy_2D is True
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
ignore_axes = (0, )
tr_transforms.append(Convert3DTo2DTransform())
else:
ignore_axes = None
tr_transforms.append(
SpatialTransform(patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"),
sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"),
angle_x=params.get("rotation_x"),
angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"),
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=order_data,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=order_seg,
random_crop=params.get("random_crop"),
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot"),
independent_scale_for_each_axis=params.get(
"independent_scale_factor_for_each_axis"),
p_independent_scale_per_axis=params.get(
"p_independent_scale_per_axis")))
if params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
# we need to put the color augmentations after the dummy 2d part (if applicable). Otherwise the overloaded color
# channel gets in the way
tr_transforms.append(GaussianNoiseTransform(p_per_sample=0.15))
tr_transforms.append(
GaussianBlurTransform((0.5, 1.5),
different_sigma_per_channel=True,
p_per_sample=0.2,
p_per_channel=0.5))
tr_transforms.append(
BrightnessMultiplicativeTransform(multiplier_range=(0.70, 1.3),
p_per_sample=0.15))
tr_transforms.append(
ContrastAugmentationTransform(contrast_range=(0.65, 1.5),
p_per_sample=0.15))
tr_transforms.append(
SimulateLowResolutionTransform(zoom_range=(0.5, 1),
per_channel=True,
p_per_channel=0.5,
order_downsample=0,
order_upsample=3,
p_per_sample=0.25,
ignore_axes=ignore_axes))
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
True,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=0.15)) # inverted gamma
# do_additive_brightness is True
if params.get("do_additive_brightness"):
tr_transforms.append(
BrightnessTransform(
params.get("additive_brightness_mu"),
params.get("additive_brightness_sigma"),
True,
p_per_sample=params.get("additive_brightness_p_per_sample"),
p_per_channel=params.get("additive_brightness_p_per_channel")))
# do_gamma is True
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
False,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
# do_mirror is True
if params.get("do_mirror") or params.get("mirror"):
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
# mask_was_used_for_normalization is OrderedDict([(0, False)]),
if params.get("mask_was_used_for_normalization") is not None:
mask_was_used_for_normalization = params.get(
"mask_was_used_for_normalization")
tr_transforms.append(
MaskTransform(mask_was_used_for_normalization,
mask_idx_in_seg=0,
set_outside_to=0))
tr_transforms.append(RemoveLabelTransform(-1, 0))
# move_last_seg_chanel_to_data is False
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
tr_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
if params.get("cascade_do_cascade_augmentations"
) and not None and params.get(
"cascade_do_cascade_augmentations"):
if params.get("cascade_random_binary_transform_p") > 0:
tr_transforms.append(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")),
0)),
p_per_sample=params.get(
"cascade_random_binary_transform_p"),
key="data",
strel_size=params.get(
"cascade_random_binary_transform_size")))
if params.get("cascade_remove_conn_comp_p") > 0:
tr_transforms.append(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")),
0)),
key="data",
p_per_sample=params.get("cascade_remove_conn_comp_p"),
fill_with_other_class_p=params.get(
"cascade_remove_conn_comp_max_size_percent_threshold"
),
dont_do_if_covers_more_than_X_percent=params.get(
"cascade_remove_conn_comp_fill_with_other_class_p")
))
tr_transforms.append(RenameTransform('seg', 'target', True))
# regions is None
if regions is not None:
tr_transforms.append(
ConvertSegmentationToRegionsTransform(regions, 'target', 'target'))
# deep_supervision_scales is a not None
if deep_supervision_scales is not None:
# soft_ds is False
if soft_ds:
assert classes is not None
tr_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'target',
'target', classes))
else:
tr_transforms.append(
DownsampleSegForDSTransform2(deep_supervision_scales,
0,
0,
input_key='target',
output_key='target'))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=seeds_train,
pin_memory=pin_memory)
# ========================================================
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
# selected_data_channels is None
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
# selected_seg_channels is [0]
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# move_last_seg_chanel_to_data is False
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
val_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
val_transforms.append(RenameTransform('seg', 'target', True))
# regions is None
if regions is not None:
val_transforms.append(
ConvertSegmentationToRegionsTransform(regions, 'target', 'target'))
# deep_supervision_scales is not None
if deep_supervision_scales is not None:
# soft_ds is False
if soft_ds:
assert classes is not None
val_transforms.append(
DownsampleSegForDSTransform3(deep_supervision_scales, 'target',
'target', classes))
else:
val_transforms.append(
DownsampleSegForDSTransform2(deep_supervision_scales,
0,
0,
input_key='target',
output_key='target'))
val_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=seeds_val,
pin_memory=pin_memory)
return batchgenerator_train, batchgenerator_val
def Transforms(patch_size,
params=default_3D_augmentation_params,
border_val_seg=-1):
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(params.get("selected_data_channels"),
data_key="data"))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert3DTo2DTransform())
tr_transforms.append(
SpatialTransform(patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"),
sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"),
angle_x=params.get("rotation_x"),
angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"),
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=3,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=1,
random_crop=params.get("random_crop"),
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot")))
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
False,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
if params.get("mask_was_used_for_normalization") is not None:
mask_was_used_for_normalization = params.get(
"mask_was_used_for_normalization")
tr_transforms.append(
MaskTransform(mask_was_used_for_normalization,
mask_idx_in_seg=0,
set_outside_to=0))
tr_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
tr_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
if params.get(
"advanced_pyramid_augmentations") and not None and params.get(
"advanced_pyramid_augmentations"):
tr_transforms.append(
ApplyRandomBinaryOperatorTransform(channel_idx=list(
range(-len(params.get("all_segmentation_labels")), 0)),
p_per_sample=0.4,
key="data",
strel_size=(1, 8)))
tr_transforms.append(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")), 0)),
key="data",
p_per_sample=0.2,
fill_with_other_class_p=0.0,
dont_do_if_covers_more_than_X_percent=0.15))
tr_transforms.append(RenameTransform('seg', 'target', True))
tr_transforms.append(NumpyToTensor(['data', 'target'], 'float'))
tr_transforms = Compose(tr_transforms)
return tr_transforms
def get_default_augmentation_withEDT(dataloader_train,
dataloader_val,
patch_size,
idx_of_edts,
params=default_3D_augmentation_params,
border_val_seg=-1,
pin_memory=True,
seeds_train=None,
seeds_val=None):
tr_transforms = []
if params.get("selected_data_channels") is not None:
tr_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
tr_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
# don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!!
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert3DTo2DTransform())
tr_transforms.append(
SpatialTransform(patch_size,
patch_center_dist_from_border=None,
do_elastic_deform=params.get("do_elastic"),
alpha=params.get("elastic_deform_alpha"),
sigma=params.get("elastic_deform_sigma"),
do_rotation=params.get("do_rotation"),
angle_x=params.get("rotation_x"),
angle_y=params.get("rotation_y"),
angle_z=params.get("rotation_z"),
do_scale=params.get("do_scaling"),
scale=params.get("scale_range"),
border_mode_data=params.get("border_mode_data"),
border_cval_data=0,
order_data=3,
border_mode_seg="constant",
border_cval_seg=border_val_seg,
order_seg=1,
random_crop=params.get("random_crop"),
p_el_per_sample=params.get("p_eldef"),
p_scale_per_sample=params.get("p_scale"),
p_rot_per_sample=params.get("p_rot")))
if params.get("dummy_2D") is not None and params.get("dummy_2D"):
tr_transforms.append(Convert2DTo3DTransform())
"""
##############################################################
##############################################################
Here we insert moving the EDT to a different key so that it does not get intensity transformed
##############################################################
##############################################################
"""
tr_transforms.append(
AppendChannelsTransform("data",
"bound",
idx_of_edts,
remove_from_input=True))
if params.get("do_gamma"):
tr_transforms.append(
GammaTransform(params.get("gamma_range"),
False,
True,
retain_stats=params.get("gamma_retain_stats"),
p_per_sample=params["p_gamma"]))
tr_transforms.append(MirrorTransform(params.get("mirror_axes")))
if params.get("mask_was_used_for_normalization") is not None:
mask_was_used_for_normalization = params.get(
"mask_was_used_for_normalization")
tr_transforms.append(
MaskTransform(mask_was_used_for_normalization,
mask_idx_in_seg=0,
set_outside_to=0))
tr_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
tr_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
if params.get(
"advanced_pyramid_augmentations") and not None and params.get(
"advanced_pyramid_augmentations"):
tr_transforms.append(
ApplyRandomBinaryOperatorTransform(channel_idx=list(
range(-len(params.get("all_segmentation_labels")), 0)),
p_per_sample=0.4,
key="data",
strel_size=(1, 8)))
tr_transforms.append(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(
range(-len(params.get("all_segmentation_labels")), 0)),
key="data",
p_per_sample=0.2,
fill_with_other_class_p=0.0,
dont_do_if_covers_more_than_X_percent=0.15))
tr_transforms.append(RenameTransform('seg', 'target', True))
tr_transforms.append(NumpyToTensor(['data', 'target', 'bound'], 'float'))
tr_transforms = Compose(tr_transforms)
batchgenerator_train = MultiThreadedAugmenter(
dataloader_train,
tr_transforms,
params.get('num_threads'),
params.get("num_cached_per_thread"),
seeds=seeds_train,
pin_memory=pin_memory)
val_transforms = []
val_transforms.append(RemoveLabelTransform(-1, 0))
if params.get("selected_data_channels") is not None:
val_transforms.append(
DataChannelSelectionTransform(
params.get("selected_data_channels")))
if params.get("selected_seg_channels") is not None:
val_transforms.append(
SegChannelSelectionTransform(params.get("selected_seg_channels")))
"""
##############################################################
##############################################################
Here we insert moving the EDT to a different key
##############################################################
##############################################################
"""
val_transforms.append(
AppendChannelsTransform("data",
"bound",
idx_of_edts,
remove_from_input=True))
if params.get("move_last_seg_chanel_to_data") is not None and params.get(
"move_last_seg_chanel_to_data"):
val_transforms.append(
MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"),
'seg', 'data'))
val_transforms.append(RenameTransform('seg', 'target', True))
val_transforms.append(NumpyToTensor(['data', 'target', 'bound'], 'float'))
val_transforms = Compose(val_transforms)
batchgenerator_val = MultiThreadedAugmenter(
dataloader_val,
val_transforms,
max(params.get('num_threads') // 2, 1),
params.get("num_cached_per_thread"),
seeds=seeds_val,
pin_memory=pin_memory)
return batchgenerator_train, batchgenerator_val
def create_data_gen_train(patient_data_train,
INPUT_PATCH_SIZE,
num_classes,
BATCH_SIZE,
contrast_range=(0.75, 1.5),
gamma_range=(0.6, 2),
num_workers=5,
num_cached_per_worker=3,
do_elastic_transform=False,
alpha=(0., 1300.),
sigma=(10., 13.),
do_rotation=False,
a_x=(0., 2 * np.pi),
a_y=(0., 2 * np.pi),
a_z=(0., 2 * np.pi),
do_scale=True,
scale_range=(0.75, 1.25),
seeds=None):
if seeds is None:
seeds = [None] * num_workers
elif seeds == 'range':
seeds = range(num_workers)
else:
assert len(seeds) == num_workers
data_gen_train = BatchGenerator3D_random_sampling(patient_data_train,
BATCH_SIZE,
num_batches=None,
seed=False,
patch_size=(160, 192,
160),
convert_labels=True)
tr_transforms = []
tr_transforms.append(DataChannelSelectionTransform([0, 1, 2, 3]))
tr_transforms.append(GenerateBrainMaskTransform())
tr_transforms.append(Mirror())
tr_transforms.append(
SpatialTransform(INPUT_PATCH_SIZE,
list(np.array(INPUT_PATCH_SIZE) // 2.),
do_elastic_deform=do_elastic_transform,
alpha=alpha,
sigma=sigma,
do_rotation=do_rotation,
angle_x=a_x,
angle_y=a_y,
angle_z=a_z,
do_scale=do_scale,
scale=scale_range,
border_mode_data='nearest',
border_cval_data=0,
order_data=3,
border_mode_seg='constant',
border_cval_seg=0,
order_seg=0,
random_crop=True))
tr_transforms.append(BrainMaskAwareStretchZeroOneTransform((-5, 5), True))
tr_transforms.append(ContrastAugmentationTransform(contrast_range, True))
tr_transforms.append(GammaTransform(gamma_range, False))
tr_transforms.append(
BrainMaskAwareStretchZeroOneTransform(per_channel=True))
tr_transforms.append(BrightnessTransform(0.0, 0.1, True))
tr_transforms.append(SegChannelSelectionTransform([0]))
tr_transforms.append(
ConvertSegToOnehotTransform(range(num_classes), 0, "seg_onehot"))
gen_train = MultiThreadedAugmenter(data_gen_train, Compose(tr_transforms),
num_workers, num_cached_per_worker,
seeds)
gen_train.restart()
return gen_train
def run(fold=0):
print fold
# =================================================================================================================
I_AM_FOLD = fold
np.random.seed(65432)
lasagne.random.set_rng(np.random.RandomState(98765))
sys.setrecursionlimit(2000)
BATCH_SIZE = 2
INPUT_PATCH_SIZE = (128, 128, 128)
num_classes = 4
EXPERIMENT_NAME = "final"
results_dir = os.path.join(paths.results_folder)
if not os.path.isdir(results_dir):
os.mkdir(results_dir)
results_dir = os.path.join(results_dir, EXPERIMENT_NAME)
if not os.path.isdir(results_dir):
os.mkdir(results_dir)
results_dir = os.path.join(results_dir, "fold%d" % I_AM_FOLD)
if not os.path.isdir(results_dir):
os.mkdir(results_dir)
n_epochs = 300
lr_decay = np.float32(0.985)
base_lr = np.float32(0.0005)
n_batches_per_epoch = 100
n_test_batches = 10
n_feedbacks_per_epoch = 10.
num_workers = 6
workers_seeds = [123, 1234, 12345, 123456, 1234567, 12345678]
# =================================================================================================================
all_data = load_dataset()
keys_sorted = np.sort(all_data.keys())
crossval_folds = KFold(len(all_data.keys()),
n_folds=5,
shuffle=True,
random_state=123456)
ctr = 0
for train_idx, test_idx in crossval_folds:
print len(train_idx), len(test_idx)
if ctr == I_AM_FOLD:
train_keys = [keys_sorted[i] for i in train_idx]
test_keys = [keys_sorted[i] for i in test_idx]
break
ctr += 1
train_data = {i: all_data[i] for i in train_keys}
test_data = {i: all_data[i] for i in test_keys}
data_gen_train = create_data_gen_train(train_data,
INPUT_PATCH_SIZE,
num_classes,
BATCH_SIZE,
contrast_range=(0.75, 1.5),
gamma_range=(0.8, 1.5),
num_workers=num_workers,
num_cached_per_worker=2,
do_elastic_transform=True,
alpha=(0., 1300.),
sigma=(10., 13.),
do_rotation=True,
a_x=(0., 2 * np.pi),
a_y=(0., 2 * np.pi),
a_z=(0., 2 * np.pi),
do_scale=True,
scale_range=(0.75, 1.25),
seeds=workers_seeds)
data_gen_validation = BatchGenerator3D_random_sampling(
test_data,
BATCH_SIZE,
num_batches=None,
seed=False,
patch_size=INPUT_PATCH_SIZE,
convert_labels=True)
val_transforms = []
val_transforms.append(GenerateBrainMaskTransform())
val_transforms.append(
BrainMaskAwareStretchZeroOneTransform(clip_range=(-5, 5),
per_channel=True))
val_transforms.append(SegChannelSelectionTransform([0]))
val_transforms.append(
ConvertSegToOnehotTransform(range(4), 0, "seg_onehot"))
val_transforms.append(DataChannelSelectionTransform([0, 1, 2, 3]))
data_gen_validation = MultiThreadedAugmenter(data_gen_validation,
Compose(val_transforms), 2, 2)
x_sym = T.tensor5()
seg_sym = T.matrix()
net, seg_layer = build_net(x_sym,
INPUT_PATCH_SIZE,
num_classes,
4,
16,
batch_size=BATCH_SIZE,
do_instance_norm=True)
output_layer_for_loss = net
# add some weight decay
l2_loss = lasagne.regularization.regularize_network_params(
output_layer_for_loss, lasagne.regularization.l2) * 1e-5
# the distinction between prediction_train and test is important only if we enable dropout (batch norm/inst norm
# does not use or save moving averages)
prediction_train = lasagne.layers.get_output(
output_layer_for_loss,
x_sym,
deterministic=False,
batch_norm_update_averages=False,
batch_norm_use_averages=False)
loss_vec = -soft_dice_per_img_in_batch(prediction_train, seg_sym,
BATCH_SIZE)[:, 1:]
loss = loss_vec.mean()
loss += l2_loss
acc_train = T.mean(T.eq(T.argmax(prediction_train, axis=1),
seg_sym.argmax(-1)),
dtype=theano.config.floatX)
prediction_test = lasagne.layers.get_output(
output_layer_for_loss,
x_sym,
deterministic=True,
batch_norm_update_averages=False,
batch_norm_use_averages=False)
loss_val = -soft_dice_per_img_in_batch(prediction_test, seg_sym,
BATCH_SIZE)[:, 1:]
loss_val = loss_val.mean()
loss_val += l2_loss
acc = T.mean(T.eq(T.argmax(prediction_test, axis=1), seg_sym.argmax(-1)),
dtype=theano.config.floatX)
# learning rate has to be a shared variable because we decrease it with every epoch
params = lasagne.layers.get_all_params(output_layer_for_loss,
trainable=True)
learning_rate = theano.shared(base_lr)
updates = lasagne.updates.adam(T.grad(loss, params),
params,
learning_rate=learning_rate,
beta1=0.9,
beta2=0.999)
dc = hard_dice_per_img_in_batch(prediction_test, seg_sym.argmax(1),
num_classes, BATCH_SIZE).mean(0)
train_fn = theano.function([x_sym, seg_sym], [loss, acc_train, loss_vec],
updates=updates)
val_fn = theano.function([x_sym, seg_sym], [loss_val, acc, dc])
all_val_dice_scores = None
all_training_losses = []
all_validation_losses = []
all_validation_accuracies = []
all_training_accuracies = []
val_dice_scores = []
epoch = 0
while epoch < n_epochs:
if epoch == 100:
data_gen_train = create_data_gen_train(train_data,
INPUT_PATCH_SIZE,
num_classes,
BATCH_SIZE,
contrast_range=(0.85, 1.25),
gamma_range=(0.8, 1.5),
num_workers=6,
num_cached_per_worker=2,
do_elastic_transform=True,
alpha=(0., 1000.),
sigma=(10., 13.),
do_rotation=True,
a_x=(0., 2 * np.pi),
a_y=(-np.pi / 8.,
np.pi / 8.),
a_z=(-np.pi / 8.,
np.pi / 8.),
do_scale=True,
scale_range=(0.85, 1.15),
seeds=workers_seeds)
if epoch == 175:
data_gen_train = create_data_gen_train(train_data,
INPUT_PATCH_SIZE,
num_classes,
BATCH_SIZE,
contrast_range=(0.9, 1.1),
gamma_range=(0.85, 1.3),
num_workers=6,
num_cached_per_worker=2,
do_elastic_transform=True,
alpha=(0., 750.),
sigma=(10., 13.),
do_rotation=True,
a_x=(0., 2 * np.pi),
a_y=(-0.00001, 0.00001),
a_z=(-0.00001, 0.00001),
do_scale=True,
scale_range=(0.85, 1.15),
seeds=workers_seeds)
epoch_start_time = time.time()
learning_rate.set_value(np.float32(base_lr * lr_decay**(epoch)))
print "epoch: ", epoch, " learning rate: ", learning_rate.get_value()
train_loss = 0
train_acc_tmp = 0
train_loss_tmp = 0
batch_ctr = 0
for data_dict in data_gen_train:
data = data_dict["data"].astype(np.float32)
seg = data_dict["seg_onehot"].astype(np.float32).transpose(
0, 2, 3, 4, 1).reshape((-1, num_classes))
if batch_ctr != 0 and batch_ctr % int(
np.floor(
n_batches_per_epoch / n_feedbacks_per_epoch)) == 0:
print "number of batches: ", batch_ctr, "/", n_batches_per_epoch
print "training_loss since last update: ", \
train_loss_tmp/np.floor(n_batches_per_epoch/n_feedbacks_per_epoch), " train accuracy: ", \
train_acc_tmp/np.floor(n_batches_per_epoch/n_feedbacks_per_epoch)
all_training_losses.append(
train_loss_tmp /
np.floor(n_batches_per_epoch / n_feedbacks_per_epoch))
all_training_accuracies.append(
train_acc_tmp /
np.floor(n_batches_per_epoch / n_feedbacks_per_epoch))
train_loss_tmp = 0
train_acc_tmp = 0
if len(val_dice_scores) > 0:
all_val_dice_scores = np.concatenate(val_dice_scores,
axis=0).reshape(
(-1, num_classes))
try:
printLosses(all_training_losses,
all_training_accuracies,
all_validation_losses,
all_validation_accuracies,
os.path.join(results_dir,
"%s.png" % EXPERIMENT_NAME),
n_feedbacks_per_epoch,
val_dice_scores=all_val_dice_scores,
val_dice_scores_labels=[
"brain", "1", "2", "3", "4", "5"
])
except:
pass
loss_vec, acc, l = train_fn(data, seg)
loss = loss_vec.mean()
train_loss += loss
train_loss_tmp += loss
train_acc_tmp += acc
batch_ctr += 1
if batch_ctr >= n_batches_per_epoch:
break
all_training_losses.append(
train_loss_tmp /
np.floor(n_batches_per_epoch / n_feedbacks_per_epoch))
all_training_accuracies.append(
train_acc_tmp /
np.floor(n_batches_per_epoch / n_feedbacks_per_epoch))
train_loss /= n_batches_per_epoch
print "training loss average on epoch: ", train_loss
val_loss = 0
accuracies = []
valid_batch_ctr = 0
all_dice = []
for data_dict in data_gen_validation:
data = data_dict["data"].astype(np.float32)
seg = data_dict["seg_onehot"].astype(np.float32).transpose(
0, 2, 3, 4, 1).reshape((-1, num_classes))
w = np.zeros(num_classes, dtype=np.float32)
w[np.unique(seg.argmax(-1))] = 1
loss, acc, dice = val_fn(data, seg)
dice[w == 0] = 2
all_dice.append(dice)
val_loss += loss
accuracies.append(acc)
valid_batch_ctr += 1
if valid_batch_ctr >= n_test_batches:
break
all_dice = np.vstack(all_dice)
dice_means = np.zeros(num_classes)
for i in range(num_classes):
dice_means[i] = all_dice[all_dice[:, i] != 2, i].mean()
val_loss /= n_test_batches
print "val loss: ", val_loss
print "val acc: ", np.mean(accuracies), "\n"
print "val dice: ", dice_means
print "This epoch took %f sec" % (time.time() - epoch_start_time)
val_dice_scores.append(dice_means)
all_validation_losses.append(val_loss)
all_validation_accuracies.append(np.mean(accuracies))
all_val_dice_scores = np.concatenate(val_dice_scores, axis=0).reshape(
(-1, num_classes))
try:
printLosses(
all_training_losses,
all_training_accuracies,
all_validation_losses,
all_validation_accuracies,
os.path.join(results_dir, "%s.png" % EXPERIMENT_NAME),
n_feedbacks_per_epoch,
val_dice_scores=all_val_dice_scores,
val_dice_scores_labels=["brain", "1", "2", "3", "4", "5"])
except:
pass
with open(
os.path.join(results_dir, "%s_Params.pkl" % (EXPERIMENT_NAME)),
'w') as f:
cPickle.dump(
lasagne.layers.get_all_param_values(output_layer_for_loss), f)
with open(
os.path.join(results_dir,
"%s_allLossesNAccur.pkl" % (EXPERIMENT_NAME)),
'w') as f:
cPickle.dump([
all_training_losses, all_training_accuracies,
all_validation_losses, all_validation_accuracies,
val_dice_scores
], f)
epoch += 1
