def random_rotate_2D_or_3D(img,
mask,
probabillity=0.8,
shift_limit=0.0625,
scale_limit=0.0,
rotate_limit=0):
"""
Rotate, shift and scale an image within a given range
:param img: numpy.ndarray
:param mask: numpy.ndarray
:param probabillity: float, will be interpreted as %-value
:param shift_limit:
:param scale_limit:
:param rotate_limit:
:return:
"""
logging.debug('random rotate for: {}'.format(img.shape))
augmented = {'image': None, 'mask': None}
if isinstance(img, sitk.Image):
img = sitk.GetArrayFromImage(img).astype(np.float32)
if isinstance(mask, sitk.Image):
mask = sitk.GetArrayFromImage(mask).astype(np.float32)
# dont print anything if no images nor masks are given
if img is None and mask is None:
logging.error('No image data given')
raise ('No image data given in grid dissortion')
# replace mask with empty slice if none is given
if mask is None:
mask = np.zeros(img.shape)
# replace image with empty slice if none is given
if img is None:
img = np.zeros(mask.shape)
if img.ndim is 2:
aug = ShiftScaleRotate(shift_limit=shift_limit,
scale_limit=scale_limit,
rotate_limit=rotate_limit,
border_mode=cv2.BORDER_REFLECT_101,
p=probabillity)
params = aug.get_params()
image_aug = aug.apply(img, interpolation=cv2.INTER_LINEAR, **params)
mask_aug = aug.apply(mask, interpolation=cv2.INTER_NEAREST, **params)
# apply shift-scale and rotation augmentation on 2d data
augmented['image'] = image_aug
augmented['mask'] = mask_aug
elif img.ndim is 3:
# apply shif-scale and rotation on 3d data, apply the same transform to all slices
images = []
masks = []
aug = ShiftScaleRotate(shift_limit=shift_limit,
scale_limit=scale_limit,
rotate_limit=rotate_limit,
border_mode=cv2.BORDER_REFLECT_101,
p=probabillity)
params = aug.get_params()
for z in range(img.shape[0]):
images.append(
aug.apply(img[z, ...],
interpolation=cv2.INTER_LINEAR,
**params))
masks.append(
aug.apply(mask[z, ...],
interpolation=cv2.INTER_NEAREST,
**params))
augmented['image'] = np.stack(images, axis=0)
augmented['mask'] = np.stack(masks, axis=0)
else:
logging.error('Unsupported dim: {}, shape: {}'.format(
img.ndim, img.shape))
raise ('Wrong shape Exception in: {}'.format('show_2D_or_3D()'))
return augmented['image'], augmented['mask']
def get_aug():
aug = ShiftScaleRotate(shift_limit=0.0325,
scale_limit=0.20,
rotate_limit=15,
p=1)
return aug.get_params()