def __call__(self, **data_dict):
data = data_dict.get(self.data_key)
seg = data_dict.get(self.label_key)
if self.patch_size is None:
if len(data.shape) == 4:
patch_size = (data.shape[2], data.shape[3])
elif len(data.shape) == 5:
patch_size = (data.shape[2], data.shape[3], data.shape[4])
else:
raise ValueError("only support 2D/3D batch data.")
else:
patch_size = self.patch_size
ret_val = augment_spatial(data, seg, patch_size=patch_size,
patch_center_dist_from_border=self.patch_center_dist_from_border,
do_elastic_deform=self.do_elastic_deform, alpha=self.alpha, sigma=self.sigma,
do_rotation=self.do_rotation, angle_x=self.angle_x, angle_y=self.angle_y,
angle_z=self.angle_z, do_scale=self.do_scale, scale=self.scale,
border_mode_data=self.border_mode_data,
border_cval_data=self.border_cval_data, order_data=self.order_data,
border_mode_seg=self.border_mode_seg, border_cval_seg=self.border_cval_seg,
order_seg=self.order_seg, random_crop=self.random_crop)
data_dict[self.data_key] = ret_val[0]
if seg is not None:
data_dict[self.label_key] = ret_val[1]
return data_dict
def __call__(self, **data_dict):
data = data_dict.get(self.data_key)
seg = data_dict.get(self.label_key)
if self.patch_size is None:
if len(data.shape) == 4:
patch_size = (data.shape[2], data.shape[3])
elif len(data.shape) == 5:
patch_size = (data.shape[2], data.shape[3], data.shape[4])
else:
raise ValueError("only support 2D/3D batch data.")
else:
patch_size = self.patch_size
ret_val = augment_spatial(data, seg, patch_size=patch_size,
patch_center_dist_from_border=self.patch_center_dist_from_border,
do_elastic_deform=self.do_elastic_deform, alpha=self.alpha, sigma=self.sigma,
do_rotation=self.do_rotation, angle_x=self.angle_x, angle_y=self.angle_y,
angle_z=self.angle_z, do_scale=self.do_scale, scale=self.scale,
border_mode_data=self.border_mode_data,
border_cval_data=self.border_cval_data, order_data=self.order_data,
border_mode_seg=self.border_mode_seg, border_cval_seg=self.border_cval_seg,
order_seg=self.order_seg, random_crop=self.random_crop)
data_dict[self.data_key] = ret_val[0]
if seg is not None:
data_dict[self.label_key] = ret_val[1]
return data_dict
def transforms(volume, segmentation, patch_size_, fraction_, spacing_,
variance_):
# BEWARE SPACING: dset.spacing = (x,y,z) but volume is shaped (z,x,y))
volume, augmentation, rval = augment_spatial(
volume,
segmentation,
patch_size=patch_size_,
patch_center_dist_from_border=np.array(patch_size_) // 2,
border_mode_data='constant',
border_mode_seg='constant',
border_cval_data=np.min(volume),
border_cval_seg=0,
alpha=(0, 750),
sigma=(10, 13),
scale=(0.8, 1.2),
do_elastic_deform=True,
do_scale=True,
do_rotation=True,
angle_x=(0, 2 * np.pi),
angle_y=(0, 0),
angle_z=(0, 0),
fraction=fraction_,
spacing=spacing_)
if np.any(variance_ != 0):
volume = augment_gaussian_noise(volume, noise_variance=variance_)
return volume, augmentation, rval
def spatial_augmentation_generator(generator, patch_size, patch_center_dist_from_border=30,
do_elastic_deform=True, alpha=(0., 1000.), sigma=(10., 13.),
do_rotation=True, angle_x=(0, 2 * np.pi), angle_y=(0, 2 * np.pi),
angle_z=(0, 2 * np.pi),
do_scale=True, scale=(0.75, 1.25), 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):
'''
THE ultimate generator. It has all you need. It alleviates the problem of having to crop your data to a reasonably sized patch size before plugging it into the
old ultimate_transform generator (In the old one you would put in patches larger than your final patch size so that rotations and deformations to not introduce black borders).
Before: Large crops = no borders but slow, small crops = black borders (duh).
Here you can just plug in the whole uncropped image and get your desired patch size as output, without performance loss or black borders
:param generator:
:param do_elastic_deform:
:param alpha:
:param sigma:
:param do_rotation:
:param angle_x:
:param angle_y:
:param angle_z:
:param do_scale:
:param scale:
:return:
'''
if not (isinstance(alpha, list) or isinstance(alpha, tuple)):
alpha = [alpha, alpha]
if not (isinstance(sigma, list) or isinstance(sigma, tuple)):
sigma = [sigma, sigma]
for data_dict in generator:
assert "data" in list(
data_dict.keys()), "your data generator needs to return a python dictionary with at least a 'data' key value pair"
data = data_dict["data"]
do_seg = False
seg = None
shape = patch_size
assert len(shape) == len(data.shape[2:]), "dimension of patch_size and data must match!"
if "seg" in list(data_dict.keys()):
seg = data_dict["seg"]
do_seg = True
data_result, seg_result = augment_spatial(data, seg, patch_size, patch_center_dist_from_border,
do_elastic_deform, alpha, sigma, do_rotation, angle_x, angle_y,
angle_z, do_scale, scale, border_mode_data, border_cval_data,
order_data, border_mode_seg, border_cval_seg, order_seg,
random_crop)
if do_seg:
data_dict['seg'] = seg_result
data_dict['data'] = data_result
yield data_dict
def test_spatial_transform(self):
img = np.zeros((1,3,10,10,10))
img[:,:, 2:4, 2:4, 2:4] = 1
lbl = np.zeros((10,10,10))
lbl[ 2:4, 2:4, 2:4] = 1
crop_size = (10,10,10)
lbl = np.expand_dims(np.expand_dims(lbl, axis=0), axis=0)
for _ in range(10):
img_ret, lbl_ret = augment_spatial(img, lbl, crop_size, patch_center_dist_from_border=30,
do_elastic_deform=False, alpha=(0., 1000.), sigma=(10., 13.),
do_rotation=False, angle_x=(0, 2 * np.pi), angle_y=(0, 2 * np.pi), angle_z=(0, 2 * np.pi),
do_scale=False, scale=(1, 1.25), border_mode_data='nearest', border_cval_data=0, order_data=3,
border_mode_seg='constant', border_cval_seg=0, order_seg=0, crop_mode='roi', p_el_per_sample=1,
p_scale_per_sample=1, p_rot_per_sample=1)
img_sum = np.sum(img)
lbl_sum = np.sum(lbl)
print(img_sum , ' ' ,np.sum(img_ret))
print(lbl_sum , ' ' ,np.sum(lbl_ret))
print(lbl_ret[0,0,2])
assert(img_sum == 3 * lbl_sum)
def __call__(self, sample):
sample = sample[None,None,:,:]
_,sample = augment_spatial(sample, sample, **self.params)
return sample[0,0]
def spatial_augmentation_generator(generator,
patch_size,
patch_center_dist_from_border=30,
do_elastic_deform=True,
alpha=(0., 1000.),
sigma=(10., 13.),
do_rotation=True,
angle_x=(0, 2 * np.pi),
angle_y=(0, 2 * np.pi),
angle_z=(0, 2 * np.pi),
do_scale=True,
scale=(0.75, 1.25),
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):
'''
THE ultimate generator. It has all you need. It alleviates the problem of having to crop your data to a reasonably sized patch size before plugging it into the
old ultimate_transform generator (In the old one you would put in patches larger than your final patch size so that rotations and deformations to not introduce black borders).
Before: Large crops = no borders but slow, small crops = black borders (duh).
Here you can just plug in the whole uncropped image and get your desired patch size as output, without performance loss or black borders
:param generator:
:param do_elastic_deform:
:param alpha:
:param sigma:
:param do_rotation:
:param angle_x:
:param angle_y:
:param angle_z:
:param do_scale:
:param scale:
:return:
'''
if not (isinstance(alpha, list) or isinstance(alpha, tuple)):
alpha = [alpha, alpha]
if not (isinstance(sigma, list) or isinstance(sigma, tuple)):
sigma = [sigma, sigma]
for data_dict in generator:
assert "data" in list(
data_dict.keys()
), "your data generator needs to return a python dictionary with at least a 'data' key value pair"
data = data_dict["data"]
do_seg = False
seg = None
shape = patch_size
assert len(shape) == len(
data.shape[2:]), "dimension of patch_size and data must match!"
if "seg" in list(data_dict.keys()):
seg = data_dict["seg"]
do_seg = True
data_result, seg_result = augment_spatial(
data, seg, patch_size, patch_center_dist_from_border,
do_elastic_deform, alpha, sigma, do_rotation, angle_x, angle_y,
angle_z, do_scale, scale, border_mode_data, border_cval_data,
order_data, border_mode_seg, border_cval_seg, order_seg,
random_crop)
if do_seg:
data_dict['seg'] = seg_result
data_dict['data'] = data_result
yield data_dict