def _pad_sample(sample, size):
padding = SegCollate._compute_padding(sample['image'].shape[1:3], size)
if padding is not None:
sample = sample.copy()
sample['image'] = SegCollate._apply_padding_to_tensor(
sample['image'], padding, value=0)
if 'labels' in sample:
sample['labels'] = SegCollate._apply_padding_to_tensor(
sample['labels'], padding, value=255)
if 'mask' in sample:
sample['mask'] = SegCollate._apply_padding_to_tensor(
sample['mask'], padding, value=255)
if 'xf_pil' in sample:
dy, dx = padding[1][0], padding[2][0]
sample['xf_pil'] = affine.cat_nx2x3(
sample['xf_pil'][None, ...],
affine.translation_matrices(np.array([[dx, dy]])))[0]
if 'xf_cv' in sample:
dy, dx = padding[1][0], padding[2][0]
sample['xf_cv'] = affine.cat_nx2x3(
affine.translation_matrices(np.array([[dx, dy]])),
sample['xf_cv'][None, ...])[0]
return sample
def test_SegImageRandomCrop():
seg_img0 = _make_test_seg_image(_make_noise_test_image())
seg_img1 = _make_test_seg_image(_make_noise_test_image())
rng = mock.Mock()
rng.uniform = mock.Mock()
img_transform = datapipe.seg_transforms_cv.SegCVTransformRandomCrop(
crop_size=(128, 128), crop_offset=(16, 16), rng=rng)
# Apply to single image
rng.uniform.side_effect = [np.array([0.5, 0.5])]
crop0 = img_transform.transform(seg_img0)
assert (crop0.image == seg_img0.image[64:192, 64:192, :]).all()
assert (crop0.labels == seg_img0.labels[64:192, 64:192]).all()
assert (crop0.mask == seg_img0.mask[64:192, 64:192]).all()
assert (crop0.xf == affine.translation_matrices(np.array([[-64,
-64]]))).all()
rng.uniform.side_effect = [np.array([0.25, 0.75])]
crop0 = img_transform.transform(seg_img0)
assert (crop0.image == seg_img0.image[32:160, 96:224, :]).all()
assert (crop0.labels == seg_img0.labels[32:160, 96:224]).all()
assert (crop0.mask == seg_img0.mask[32:160, 96:224]).all()
assert (crop0.xf == affine.translation_matrices(np.array([[-96,
-32]]))).all()
# Apply to paired image
# First crop at 64,64 (0.5*128), with the second at 56,72
rng.uniform.side_effect = [np.array([0.5, 0.5]), np.array([-0.5, 0.5])]
crop0, crop1 = img_transform.transform_pair(seg_img0, seg_img1)
assert (crop0.image == seg_img0.image[64:192, 64:192, :]).all()
assert (crop0.labels == seg_img0.labels[64:192, 64:192]).all()
assert (crop0.mask == seg_img0.mask[64:192, 64:192]).all()
assert (crop0.xf == affine.translation_matrices(np.array([[-64,
-64]]))).all()
assert (crop1.image == seg_img0.image[56:184, 72:200, :]).all()
assert (crop1.labels == seg_img0.labels[56:184, 72:200]).all()
assert (crop1.mask == seg_img0.mask[56:184, 72:200]).all()
assert (crop1.xf == affine.translation_matrices(np.array([[-72,
-56]]))).all()
rng.uniform.side_effect = [np.array([0.5, 0.5]), np.array([-0.5, 0.5])]
img_transform = datapipe.seg_transforms_cv.SegCVTransformRandomCrop(
crop_size=(128, 128),
crop_offset=(16, 16),
rng=np.random.RandomState(12345))
for _ in range(REPEATS):
_paired_transform_ocv_test(seg_img0.image, img_transform, (128, 128),
(1, 1))
_paired_transform_torch_test(seg_img0.image, img_transform, (128, 128),
(1, 1))
def transform_single(self, sample0):
sample0 = sample0.copy()
scale_dim = 1 if self.uniform_scale else 2
# Draw scale factor
f_scale = 0.5 + self.rng.randint(0, 11, size=(scale_dim, )) / 10.0
# Scale the crop size by the inverse of the scale
sc_size = np.round(self.crop_size_arr / f_scale).astype(int)
sample0 = self.pad_single(sample0, sc_size)
image, labels, mask, xf = sample0['image_arr'], sample0.get(
'labels_arr'), sample0.get('mask_arr'), sample0.get('xf_cv')
# Randomly choose position
extra = np.array(image.shape[:2]) - sc_size
pos = np.round(extra *
self.rng.uniform(0.0, 1.0, size=(2, ))).astype(int)
# Extract crop and scale to target size
image = image[pos[0]:pos[0] + sc_size[0], pos[1]:pos[1] + sc_size[1]]
sample0['image_arr'] = cv2.resize(image,
self.crop_size[::-1],
interpolation=cv2.INTER_LINEAR)
if labels is not None:
labels = labels[pos[0]:pos[0] + sc_size[0],
pos[1]:pos[1] + sc_size[1]]
sample0['labels_arr'] = cv2.resize(labels,
self.crop_size[::-1],
interpolation=cv2.INTER_NEAREST)
if mask is not None:
mask = mask[pos[0]:pos[0] + sc_size[0], pos[1]:pos[1] + sc_size[1]]
sample0['mask_arr'] = cv2.resize(mask,
self.crop_size[::-1],
interpolation=cv2.INTER_LINEAR)
if xf is not None:
# Matching `cv2.resize` requires:
# - scale factor of out_size/in_size
# - a translation of (scale_factor - 1) / 2
scale_factor_yx = self.crop_size_arr / sc_size
resize_xlat_yx = (scale_factor_yx - 1.0) * 0.5
sample0['xf_cv'] = affine.cat_nx2x3(
affine.translation_matrices(
resize_xlat_yx[None, ::-1].astype(float)),
affine.scale_matrices(scale_factor_yx[None, ::-1]),
affine.translation_matrices(-pos[None, ::-1].astype(float)),
xf[None, ...])[0]
return sample0
def transform_pair(self, sample0, sample1):
# Pad the image if necessary
sample0, sample1 = self.pad_pair(sample0, sample1, self.crop_size)
# Randomly choose positions of each crop
extra = np.array(sample0['image_arr'].shape[:2]) - self.crop_size
pos0 = np.round(extra *
self.rng.uniform(0.0, 1.0, size=(2, ))).astype(int)
pos1 = pos0 + np.round(self.crop_offset * self.rng.uniform(
-1.0, 1.0, size=(2, ))).astype(int)
# Ensure pos1 cannot go out of bounds
pos1 = np.clip(pos1, np.array([0, 0]), extra)
# Extract crop and scale to target size
sample0['image_arr'] = sample0['image_arr'][pos0[0]:pos0[0] +
self.crop_size[0],
pos0[1]:pos0[1] +
self.crop_size[1]]
sample1['image_arr'] = sample1['image_arr'][pos1[0]:pos1[0] +
self.crop_size[0],
pos1[1]:pos1[1] +
self.crop_size[1]]
sample0['mask_arr'] = sample0['mask_arr'][pos0[0]:pos0[0] +
self.crop_size[0],
pos0[1]:pos0[1] +
self.crop_size[1]]
sample1['mask_arr'] = sample1['mask_arr'][pos1[0]:pos1[0] +
self.crop_size[0],
pos1[1]:pos1[1] +
self.crop_size[1]]
if 'labels_arr' in sample0:
sample0['labels_arr'] = sample0['labels_arr'][pos0[0]:pos0[0] +
self.crop_size[0],
pos0[1]:pos0[1] +
self.crop_size[1]]
sample1['labels_arr'] = sample1['labels_arr'][pos1[0]:pos1[0] +
self.crop_size[0],
pos1[1]:pos1[1] +
self.crop_size[1]]
if 'xf_cv' in sample0:
sample0['xf_cv'] = affine.cat_nx2x3(
affine.translation_matrices(-pos0[None, ::-1]),
sample0['xf_cv'][None, ...])[0]
sample1['xf_cv'] = affine.cat_nx2x3(
affine.translation_matrices(-pos1[None, ::-1]),
sample1['xf_cv'][None, ...])[0]
return (sample0, sample1)
def test_SegDataPipeline():
seg_img0 = _make_test_seg_image(_make_noise_test_image())
rng = mock.Mock()
rng.uniform = mock.Mock()
img_transform = seg_transforms.SegTransformCompose([
datapipe.seg_transforms_cv.SegCVTransformRandomCrop(crop_size=(128,
128),
crop_offset=(16,
16),
rng=rng),
datapipe.seg_transforms_cv.SegCVTransformNormalizeToTensor(None, None),
])
pipe = seg_transforms.SegDataPipeline((32, 32), img_transform)
# Supervised batch
rng.uniform.side_effect = [np.array([0.5, 0.5])]
xs, ys = pipe.prepare_supervised_batch([seg_img0.image], [seg_img0.labels])
assert (xs == img_as_float(seg_img0.image).transpose(
2, 0, 1)[None, :, 64:192, 64:192].astype(np.float32)).all()
assert (ys == seg_img0.labels[None, None, 64:192, 64:192]).all()
# Unsupervised batch
rng.uniform.side_effect = [np.array([0.5, 0.5])]
xs, ms = pipe.prepare_unsupervised_batch([seg_img0.image])
assert (xs == img_as_float(seg_img0.image).transpose(
2, 0, 1)[None, :, 64:192, 64:192].astype(np.float32)).all()
assert (ms == np.ones((1, 1, 128, 128), dtype=np.float32)).all()
# Unsupervised paired batch
rng.uniform.side_effect = [np.array([0.5, 0.5]), np.array([-0.5, 0.5])]
x0s, m0s, xf0s, x1s, m1s, xf1s = pipe.prepare_unsupervised_paired_batch(
[seg_img0.image])
assert (x0s == img_as_float(seg_img0.image).transpose(
2, 0, 1)[None, :, 64:192, 64:192].astype(np.float32)).all()
assert (x1s == img_as_float(seg_img0.image).transpose(
2, 0, 1)[None, :, 56:184, 72:200].astype(np.float32)).all()
assert (m0s == np.ones((1, 1, 128, 128), dtype=np.float32)).all()
assert (m1s == np.ones((1, 1, 128, 128), dtype=np.float32)).all()
assert (xf0s == affine.translation_matrices(np.array([[-64, -64]]))).all()
assert (xf1s == affine.translation_matrices(np.array([[-72, -56]]))).all()
def pad_pair(self, sample0, sample1, min_size):
sample0 = sample0.copy()
sample1 = sample1.copy()
image0 = sample0['image_arr']
image1 = sample1['image_arr']
img_size0 = image0.shape[:2]
if img_size0[0] < min_size[0] or img_size0[1] < min_size[1]:
# Padding required
# Compute padding
pad_h = max(min_size[0] - img_size0[0], 0)
pad_w = max(min_size[1] - img_size0[1], 0)
h0 = pad_h // 2
h1 = pad_h - h0
w0 = pad_w // 2
w1 = pad_w - w0
# Add an alpha channel to the image so that we can use it during standardisation
# to ensure that the padding area has a value of 0, post mean-subtraction
alpha_channel = np.ones(img_size0 +
(1, ), dtype=image0.dtype) * 255
image0 = np.append(image0[:, :, :3], alpha_channel, axis=2)
image1 = np.append(image1[:, :, :3], alpha_channel, axis=2)
# Apply
sample0['image_arr'] = np.pad(image0, [[h0, h1], [w0, w1], [0, 0]],
mode='constant',
constant_values=0)
sample1['image_arr'] = np.pad(image1, [[h0, h1], [w0, w1], [0, 0]],
mode='constant',
constant_values=0)
if 'labels_arr' in sample0:
sample0['labels_arr'] = np.pad(sample0['labels_arr'],
[[h0, h1], [w0, w1]],
mode='constant',
constant_values=255)
sample1['labels_arr'] = np.pad(sample1['labels_arr'],
[[h0, h1], [w0, w1]],
mode='constant',
constant_values=255)
if 'mask_arr' in sample0:
sample0['mask_arr'] = np.pad(sample0['mask_arr'],
[[h0, h1], [w0, w1]],
mode='constant')
sample1['mask_arr'] = np.pad(sample1['mask_arr'],
[[h0, h1], [w0, w1]],
mode='constant')
if 'xf_cv' in sample0:
pad_xlat = affine.translation_matrices(np.array([[w0, h0]]))
sample0['xf_cv'] = affine.cat_nx2x3(
pad_xlat, sample0['xf_cv'][None, ...])[0]
sample1['xf_cv'] = affine.cat_nx2x3(
pad_xlat, sample1['xf_cv'][None, ...])[0]
return (sample0, sample1)
def transform_pair(self, sample0, sample1):
sample0 = sample0.copy()
sample1 = sample1.copy()
# Flip flags
flip_flags_xyd = self.rng.binomial(1, 0.5, size=(2, 3)) != 0
flip_flags_xyd = flip_flags_xyd & np.array(
[[self.hflip, self.vflip, self.hvflip]])
sample0['image_arr'] = self.flip_image(sample0['image_arr'],
flip_flags_xyd[0])
sample1['image_arr'] = self.flip_image(sample1['image_arr'],
flip_flags_xyd[1])
if 'mask_arr' in sample0:
sample0['mask_arr'] = self.flip_image(sample0['mask_arr'],
flip_flags_xyd[0])
sample1['mask_arr'] = self.flip_image(sample1['mask_arr'],
flip_flags_xyd[1])
if 'labels_arr' in sample0:
sample0['labels_arr'] = self.flip_image(sample0['labels_arr'],
flip_flags_xyd[0])
sample1['labels_arr'] = self.flip_image(sample1['labels_arr'],
flip_flags_xyd[1])
if 'xf_cv' in sample0:
# False -> 1, True -> -1
flip_scale_xy = flip_flags_xyd[:, :2] * -2 + 1
# Negative scale factors need to be combined with a translation whose value is (image_size - 1)
# Mask the translation with the flip flags to only apply it where flipping is done
flip_xlat_xy = flip_flags_xyd[:, :2] * (np.array([
sample0['image_arr'].shape[:2][::-1],
sample1['image_arr'].shape[:2][::-1]
]).astype(float) - 1)
hv_flip_xf = affine.identity_xf(2)
hv_flip_xf[flip_flags_xyd[:, 2]] = hv_flip_xf[
flip_flags_xyd[:, 2], ::-1, :]
xf01 = np.stack([sample0['xf_cv'], sample1['xf_cv']], axis=0)
xf01 = affine.cat_nx2x3(
hv_flip_xf,
affine.translation_matrices(flip_xlat_xy),
affine.scale_matrices(flip_scale_xy),
xf01,
)
sample0['xf_cv'] = xf01[0]
sample1['xf_cv'] = xf01[1]
return (sample0, sample1)
def pad_single(self, sample, min_size):
sample = sample.copy()
image = sample['image_arr']
# image, labels, mask, xf = seg_img.image, seg_img.labels, seg_img.mask, seg_img.xf
img_size = image.shape[:2]
if img_size[0] < min_size[0] or img_size[1] < min_size[1]:
# Padding required
# Compute padding
pad_h = max(min_size[0] - img_size[0], 0)
pad_w = max(min_size[1] - img_size[1], 0)
h0 = pad_h // 2
h1 = pad_h - h0
w0 = pad_w // 2
w1 = pad_w - w0
# Add an alpha channel to the image so that we can use it during standardisation
# to ensure that the padding area has a value of 0, post mean-subtraction
alpha_channel = np.ones(img_size + (1, ), dtype=image.dtype) * 255
image = np.append(image[:, :, :3], alpha_channel, axis=2)
# Apply
sample['image_arr'] = np.pad(image, [[h0, h1], [w0, w1], [0, 0]],
mode='constant',
constant_values=0)
if 'labels_arr' in sample:
sample['labels_arr'] = np.pad(sample['labels_arr'],
[[h0, h1], [w0, w1]],
mode='constant',
constant_values=255)
if 'mask_arr' in sample:
sample['mask_arr'] = np.pad(sample['mask_arr'],
[[h0, h1], [w0, w1]],
mode='constant')
if 'xf_cv' in sample:
sample['xf_cv'] = affine.cat_nx2x3(
affine.translation_matrices(np.array([[w0, h0]])),
sample['xf_cv'][None, ...])[0]
return sample
def transform_single(self, sample):
sample = self.pad_single(sample, self.crop_size)
image = sample['image_arr']
extra = np.array(image.shape[:2]) - self.crop_size
pos = np.round(extra *
self.rng.uniform(0.0, 1.0, size=(2, ))).astype(int)
sample['image_arr'] = image[pos[0]:pos[0] + self.crop_size[0],
pos[1]:pos[1] + self.crop_size[1]]
if 'labels_arr' in sample:
sample['labels_arr'] = sample['labels_arr'][pos[0]:pos[0] +
self.crop_size[0],
pos[1]:pos[1] +
self.crop_size[1]]
if 'mask_arr' in sample:
sample['mask_arr'] = sample['mask_arr'][pos[0]:pos[0] +
self.crop_size[0],
pos[1]:pos[1] +
self.crop_size[1]]
if 'xf_cv' in sample:
sample['xf_cv'] = affine.cat_nx2x3(
affine.translation_matrices(-pos[None, ::-1].astype(float)),
sample['xf_cv'][None, ...])[0]
return sample
def test_SegImageRandomCropScaleHung():
seg_img0 = _make_test_seg_image(_make_noise_test_image())
seg_img1 = _make_test_seg_image(_make_noise_test_image())
rng = mock.Mock()
rng.randint = mock.Mock()
rng.uniform = mock.Mock()
img_transform = datapipe.seg_transforms_cv.SegCVTransformRandomCropScaleHung(
crop_size=(128, 128), crop_offset=(16, 16), rng=rng)
# Apply to single image
rng.randint.side_effect = [
3
] # 3 / 10 + 0.5 = scale factor of 0.8; crop of 160x160 scaled to 128x128
rng.uniform.side_effect = [np.array([0.5, 0.5])]
crop0 = img_transform.transform(seg_img0)
assert (crop0.image == cv2.resize(seg_img0.image[48:208, 48:208, :],
(128, 128),
interpolation=cv2.INTER_LINEAR)).all()
assert (crop0.labels == cv2.resize(seg_img0.labels[48:208, 48:208],
(128, 128),
interpolation=cv2.INTER_NEAREST)).all()
assert (crop0.mask == cv2.resize(seg_img0.mask[48:208, 48:208], (128, 128),
interpolation=cv2.INTER_LINEAR)).all()
assert np.allclose(crop0.xf,
np.array([[[0.8, 0.0, -38.4], [0.0, 0.8, -38.4]]]))
rng.randint.side_effect = [
7
] # 7 / 10 + 0.5 = scale factor of 1.2; crop of 107x107 scaled to 128x128
rng.uniform.side_effect = [np.array([0.25, 0.75])]
crop0 = img_transform.transform(seg_img0)
assert (crop0.image == cv2.resize(seg_img0.image[37:144, 112:219, :],
(128, 128),
interpolation=cv2.INTER_LINEAR)).all()
assert (crop0.labels == cv2.resize(seg_img0.labels[37:144, 112:219],
(128, 128),
interpolation=cv2.INTER_NEAREST)).all()
assert (crop0.mask == cv2.resize(seg_img0.mask[37:144, 112:219],
(128, 128),
interpolation=cv2.INTER_LINEAR)).all()
assert np.allclose(
crop0.xf,
np.array([[[128 / 107, 0.0, -112 * 128 / 107],
[0.0, 128 / 107, -37 * 128 / 107]]]))
# Apply to paired image
# First crop at 64,64 (0.5*128), with the second at 56,72
rng.randint.side_effect = [
3
] # 3 / 10 + 0.5 = scale factor of 0.8; crop of 160x160 scaled to 128x128
rng.uniform.side_effect = [np.array([0.5, 0.5]), np.array([-0.5, 0.5])]
# extra=(96, 96), pos0=(48,48), pos1=(40,56), centre0=(128,128), centre1=(120,136), pos0=(64,64), pos1=(40,56)
crop0, crop1 = img_transform.transform_pair(seg_img0, seg_img1)
assert (crop0.image == seg_img0.image[64:192, 64:192, :]).all()
assert (crop0.labels == seg_img0.labels[64:192, 64:192]).all()
assert (crop0.mask == seg_img0.mask[64:192, 64:192]).all()
assert (crop0.xf == affine.translation_matrices(np.array([[-64,
-64]]))).all()
assert (crop1.image == cv2.resize(seg_img1.image[40:200, 56:216, :],
(128, 128),
interpolation=cv2.INTER_LINEAR)).all()
assert (crop1.labels == cv2.resize(seg_img1.labels[40:200, 56:216],
(128, 128),
interpolation=cv2.INTER_NEAREST)).all()
assert (crop1.mask == cv2.resize(seg_img1.mask[40:200, 56:216], (128, 128),
interpolation=cv2.INTER_LINEAR)).all()
assert np.allclose(crop1.xf,
np.array([[[0.8, 0.0, -44.8], [0.0, 0.8, -32.0]]]))
def test_SegImageTransformPad_single():
"""Test SegImageTransformPad that pads an image.
"""
seg_img = _make_test_seg_image(_make_noise_test_image())
img_transform = datapipe.seg_transforms_cv.SegCVTransformPad()
# No padding required
seg_img_padded = img_transform.pad_single(seg_img, (128, 128))
assert seg_img_padded.image.shape == seg_img.image.shape
assert (seg_img_padded.image == seg_img.image).all()
assert (seg_img_padded.labels == seg_img.labels).all()
assert (seg_img_padded.mask == seg_img.mask).all()
assert (seg_img_padded.xf == affine.identity_xf(1)).all()
# 384,256, pad with 64 in y
seg_img_padded = img_transform.pad_single(seg_img, (384, 256))
assert seg_img_padded.image.shape[:2] == (384, 256)
assert (seg_img_padded.image[:, :, :3] == np.pad(
seg_img.image, [[64, 64], [0, 0], [0, 0]], mode='constant')).all()
alpha = np.full((256, 256), 255)
assert (seg_img_padded.image[:, :, 3] == np.pad(alpha, [[64, 64], [0, 0]],
mode='constant')).all()
assert (seg_img_padded.labels == np.pad(seg_img.labels, [[64, 64], [0, 0]],
mode='constant',
constant_values=255)).all()
assert (seg_img_padded.mask == np.pad(seg_img.mask, [[64, 64], [0, 0]],
mode='constant',
constant_values=0)).all()
assert (seg_img_padded.xf == affine.translation_matrices(
np.array([[0, 64]]))).all()
# 385,256, pad with 65,64 in y
seg_img_padded = img_transform.pad_single(seg_img, (385, 256))
assert seg_img_padded.image.shape[:2] == (385, 256)
assert (seg_img_padded.image[:, :, :3] == np.pad(
seg_img.image, [[64, 65], [0, 0], [0, 0]], mode='constant')).all()
alpha = np.full((256, 256), 255)
assert (seg_img_padded.image[:, :, 3] == np.pad(alpha, [[64, 65], [0, 0]],
mode='constant')).all()
assert (seg_img_padded.labels == np.pad(seg_img.labels, [[64, 65], [0, 0]],
mode='constant',
constant_values=255)).all()
assert (seg_img_padded.mask == np.pad(seg_img.mask, [[64, 65], [0, 0]],
mode='constant',
constant_values=0)).all()
assert (seg_img_padded.xf == affine.translation_matrices(
np.array([[0, 64]]))).all()
# 256,384, pad with 64 in x
seg_img_padded = img_transform.pad_single(seg_img, (256, 384))
assert seg_img_padded.image.shape[:2] == (256, 384)
assert (seg_img_padded.image[:, :, :3] == np.pad(
seg_img.image, [[0, 0], [64, 64], [0, 0]], mode='constant')).all()
assert (seg_img_padded.image[:, :, 3] == np.pad(alpha, [[0, 0], [64, 64]],
mode='constant')).all()
assert (seg_img_padded.labels == np.pad(seg_img.labels, [[0, 0], [64, 64]],
mode='constant',
constant_values=255)).all()
assert (seg_img_padded.mask == np.pad(seg_img.mask, [[0, 0], [64, 64]],
mode='constant',
constant_values=0)).all()
assert (seg_img_padded.xf == affine.translation_matrices(
np.array([[64, 0]]))).all()
# 256,385, pad with 65,64 in x
seg_img_padded = img_transform.pad_single(seg_img, (256, 385))
assert seg_img_padded.image.shape[:2] == (256, 385)
assert (seg_img_padded.image[:, :, :3] == np.pad(
seg_img.image, [[0, 0], [64, 65], [0, 0]], mode='constant')).all()
assert (seg_img_padded.image[:, :, 3] == np.pad(alpha, [[0, 0], [64, 65]],
mode='constant')).all()
assert (seg_img_padded.labels == np.pad(seg_img.labels, [[0, 0], [64, 65]],
mode='constant',
constant_values=255)).all()
assert (seg_img_padded.mask == np.pad(seg_img.mask, [[0, 0], [64, 65]],
mode='constant',
constant_values=0)).all()
assert (seg_img_padded.xf == affine.translation_matrices(
np.array([[64, 0]]))).all()
def transform_pair(self, sample0, sample1):
sample0 = sample0.copy()
sample1 = sample1.copy()
# Choose scales and rotations
if self.constrain_rot_scale:
if self.uniform_scale:
scale_factors_yx = np.exp(
self.rng.uniform(-self.log_max_scale,
self.log_max_scale,
size=(1, 1)))
scale_factors_yx = np.repeat(scale_factors_yx, 2, axis=1)
else:
scale_factors_yx = np.exp(
self.rng.uniform(-self.log_max_scale,
self.log_max_scale,
size=(1, 2)))
rot_thetas = self.rng.uniform(-self.rot_mag_rad,
self.rot_mag_rad,
size=(1, ))
scale_factors_yx = np.repeat(scale_factors_yx, 2, axis=0)
rot_thetas = np.repeat(rot_thetas, 2, axis=0)
else:
if self.uniform_scale:
scale_factors_yx = np.exp(
self.rng.uniform(-self.log_max_scale,
self.log_max_scale,
size=(2, 1)))
scale_factors_yx = np.repeat(scale_factors_yx, 2, axis=1)
else:
scale_factors_yx = np.exp(
self.rng.uniform(-self.log_max_scale,
self.log_max_scale,
size=(2, 2)))
rot_thetas = self.rng.uniform(-self.rot_mag_rad,
self.rot_mag_rad,
size=(2, ))
img_size = np.array(sample0['image_arr'].shape[:2])
# Scale the crop size by the inverse of the scale
sc_size = self.crop_size_arr / scale_factors_yx.min(axis=0)
crop_centre_pos = np.minimum(sc_size, img_size) * 0.5
# Randomly choose centres
extra = np.maximum(img_size - sc_size, 0.0)
centre0 = extra * self.rng.uniform(0.0, 1.0,
size=(2, )) + crop_centre_pos
offset1 = np.round(self.crop_offset *
self.rng.uniform(-1.0, 1.0, size=(2, )))
centre_xlat = np.stack([centre0, centre0], axis=0)
offset1_xlat = np.stack([np.zeros((2, )), offset1], axis=0)
# Build affine transformation matrices
local_xfs = affine.cat_nx2x3(
affine.translation_matrices(self.crop_size_arr[None, ::-1] * 0.5),
affine.translation_matrices(offset1_xlat[:, ::-1]),
affine.rotation_matrices(rot_thetas),
affine.scale_matrices(scale_factors_yx[:, ::-1]),
affine.translation_matrices(-centre_xlat[:, ::-1]),
)
# Use nearest neighbour sampling to stay consistent with labels, if labels present
interpolation = cv2.INTER_NEAREST if 'labels_arr' in sample0 else cv2.INTER_LINEAR
sample0['image_arr'] = cv2.warpAffine(
sample0['image_arr'],
local_xfs[0],
self.crop_size[::-1],
flags=interpolation,
borderValue=0,
borderMode=cv2.BORDER_REFLECT_101)
sample1['image_arr'] = cv2.warpAffine(
sample1['image_arr'],
local_xfs[1],
self.crop_size[::-1],
flags=interpolation,
borderValue=0,
borderMode=cv2.BORDER_REFLECT_101)
if 'labels_arr' in sample0:
sample0['labels_arr'] = cv2.warpAffine(
sample0['labels_arr'],
local_xfs[0],
self.crop_size[::-1],
flags=cv2.INTER_NEAREST,
borderValue=255,
borderMode=cv2.BORDER_CONSTANT)
sample1['labels_arr'] = cv2.warpAffine(
sample1['labels_arr'],
local_xfs[1],
self.crop_size[::-1],
flags=cv2.INTER_NEAREST,
borderValue=255,
borderMode=cv2.BORDER_CONSTANT)
if 'mask_arr' in sample0:
sample0['mask_arr'] = cv2.warpAffine(
sample0['mask_arr'],
local_xfs[0],
self.crop_size[::-1],
flags=interpolation,
borderValue=0,
borderMode=cv2.BORDER_CONSTANT)
sample1['mask_arr'] = cv2.warpAffine(
sample1['mask_arr'],
local_xfs[1],
self.crop_size[::-1],
flags=interpolation,
borderValue=0,
borderMode=cv2.BORDER_CONSTANT)
if 'xf_cv' in sample0:
xf01 = affine.cat_nx2x3(
local_xfs,
np.stack([sample0['xf_cv'], sample1['xf_cv']], axis=0))
sample0['xf_cv'] = xf01[0]
sample1['xf_cv'] = xf01[1]
return (sample0, sample1)
def transform_single(self, sample0):
sample0 = sample0.copy()
# Extract contents
image = sample0['image_arr']
# Choose scale and rotation
if self.uniform_scale:
scale_factor_yx = np.exp(
self.rng.uniform(-self.log_max_scale,
self.log_max_scale,
size=(1, )))
scale_factor_yx = np.repeat(scale_factor_yx, 2, axis=0)
else:
scale_factor_yx = np.exp(
self.rng.uniform(-self.log_max_scale,
self.log_max_scale,
size=(2, )))
rot_theta = self.rng.uniform(-self.rot_mag_rad,
self.rot_mag_rad,
size=(1, ))
# Scale the crop size by the inverse of the scale
sc_size = self.crop_size_arr / scale_factor_yx
# Randomly choose centre
img_size = np.array(image.shape[:2])
extra = np.maximum(img_size - sc_size, 0.0)
centre = extra * self.rng.uniform(0.0, 1.0, size=(2, )) + np.minimum(
sc_size, img_size) * 0.5
# Build affine transformation matrix
local_xf = affine.cat_nx2x3(
affine.translation_matrices(self.crop_size_arr[None, ::-1] * 0.5),
affine.rotation_matrices(rot_theta),
affine.scale_matrices(scale_factor_yx[None, ::-1]),
affine.translation_matrices(-centre[None, ::-1]),
)
# Reflect the image
# Use nearest neighbour sampling to stay consistent with labels, if labels present
if 'labels_arr' in sample0:
interpolation = cv2.INTER_NEAREST
else:
interpolation = self.rng.choice(
[cv2.INTER_NEAREST, cv2.INTER_LINEAR])
sample0['image_arr'] = cv2.warpAffine(
image,
local_xf[0],
self.crop_size[::-1],
flags=interpolation,
borderValue=0,
borderMode=cv2.BORDER_REFLECT_101)
# Don't reflect labels and mask
if 'labels_arr' in sample0:
sample0['labels_arr'] = cv2.warpAffine(
sample0['labels_arr'],
local_xf[0],
self.crop_size[::-1],
flags=cv2.INTER_NEAREST,
borderValue=255,
borderMode=cv2.BORDER_CONSTANT)
if 'mask_arr' in sample0:
sample0['mask_arr'] = cv2.warpAffine(
sample0['mask_arr'],
local_xf[0],
self.crop_size[::-1],
flags=interpolation,
borderValue=0,
borderMode=cv2.BORDER_CONSTANT)
if 'xf_cv' in sample0:
sample0['xf_cv'] = affine.cat_nx2x3(local_xf,
sample0['xf_cv'][None, ...])[0]
return sample0
def transform_pair(self, sample0, sample1):
sample0 = sample0.copy()
sample1 = sample1.copy()
scale_dim = 1 if self.uniform_scale else 2
# Draw a scale factor for the second crop (crop1 from crop0,crop1)
f_scale1 = 0.5 + self.rng.randint(0, 11, size=(scale_dim, )) / 10.0
# Scale the crop size by the inverse of the scale
sc_size1 = np.round(self.crop_size_arr / f_scale1).astype(int)
# Compute the maximum crop size that we need
max_sc_size = np.maximum(self.crop_size_arr, sc_size1)
# Pad the image if necessary
sample0, sample1 = self.pad_pair(sample0, sample1, max_sc_size)
# Randomly choose positions of each crop
extra = np.array(sample0['image_arr'].shape[:2]) - max_sc_size
pos0 = np.round(extra *
self.rng.uniform(0.0, 1.0, size=(2, ))).astype(int)
pos1 = pos0 + np.round(self.crop_offset * self.rng.uniform(
-1.0, 1.0, size=(2, ))).astype(int)
# Ensure pos1 cannot go out of bounds
pos1 = np.clip(pos1, np.array([0, 0]), extra)
centre0 = pos0 + max_sc_size * 0.5
centre1 = pos1 + max_sc_size * 0.5
pos0 = np.round(centre0 - self.crop_size_arr * 0.5).astype(int)
pos1 = np.round(centre1 - sc_size1 * 0.5).astype(int)
# Extract crop and scale to target size
sample0['image_arr'] = sample0['image_arr'][pos0[0]:pos0[0] +
self.crop_size_arr[0],
pos0[1]:pos0[1] +
self.crop_size_arr[1]]
sample1['image_arr'] = sample1['image_arr'][pos1[0]:pos1[0] +
sc_size1[0],
pos1[1]:pos1[1] +
sc_size1[1]]
# image0 = cv2.resize(image0, self.crop_size[::-1], interpolation=cv2.INTER_LINEAR)
sample1['image_arr'] = cv2.resize(sample1['image_arr'],
self.crop_size[::-1],
interpolation=cv2.INTER_LINEAR)
if 'mask_arr' in sample0:
sample0['mask_arr'] = sample0['mask_arr'][pos0[0]:pos0[0] +
self.crop_size_arr[0],
pos0[1]:pos0[1] +
self.crop_size_arr[1]]
sample1['mask_arr'] = sample1['mask_arr'][pos1[0]:pos1[0] +
sc_size1[0],
pos1[1]:pos1[1] +
sc_size1[1]]
# mask0 = cv2.resize(mask0, self.crop_size[::-1], interpolation=cv2.INTER_NEAREST)
sample1['mask_arr'] = cv2.resize(sample1['mask_arr'],
self.crop_size[::-1],
interpolation=cv2.INTER_NEAREST)
if 'labels_arr' in sample0:
sample0['labels_arr'] = sample0['labels_arr'][
pos0[0]:pos0[0] + self.crop_size_arr[0],
pos0[1]:pos0[1] + self.crop_size_arr[1]]
sample1['labels_arr'] = sample1['labels_arr'][pos1[0]:pos1[0] +
sc_size1[0],
pos1[1]:pos1[1] +
sc_size1[1]]
# labels0 = cv2.resize(labels0, self.crop_size[::-1], interpolation = cv2.INTER_NEAREST)
sample1['labels_arr'] = cv2.resize(sample1['labels_arr'],
self.crop_size[::-1],
interpolation=cv2.INTER_NEAREST)
if 'xf_cv' in sample0:
xf01 = np.stack([sample0['xf_cv'], sample1['xf_cv']], axis=0)
positions_xy = np.append(pos0[None, ::-1],
pos1[None, ::-1],
axis=0)
# Matching `cv2.resize` requires:
# - scale factor of out_size/in_size
# - a translation of (scale_factor - 1) / 2
scale_factors_xy = np.append(
np.array([[1, 1]]),
self.crop_size_arr[None, ::-1].astype(float) /
sc_size1[None, ::-1],
axis=0)
resize_xlats_xy = (scale_factors_xy - 1.0) * 0.5
xf01 = affine.cat_nx2x3(
affine.translation_matrices(resize_xlats_xy),
affine.scale_matrices(scale_factors_xy),
affine.translation_matrices(-positions_xy),
xf01,
)
sample0['xf_cv'] = xf01[0]
sample1['xf_cv'] = xf01[1]
return sample0, sample1