def augment_batch_img(batch_img, batch_pts, plot=False):
"""
Image augmentation, used when training
:param batch_img: [B,H,W,C]
:param batch_pts: [B,number,xy]
:return: aug_b_img, aug_b_pts
"""
sometimes = lambda aug: iaa.Sometimes(0.5, aug)
seq = iaa.Sequential([
iaa.CropAndPad(percent=((0., 0.), (-0.1, 0.1), (0., 0.), (-0.1, 0.1))),
iaa.Affine(rotate=(-10, 10)),
iaa.Add((-25, 25)) # change brightness
])
aug_b_imgs, aug_b_pts = seq(images=batch_img, keypoints=batch_pts)
if plot:
import cv2
batch_img = [
cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in batch_img
]
aug_b_imgs = [
cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in aug_b_imgs
]
for i in range(len(batch_img)):
print("[Image #%d]" % (i, ))
keypoints_before = KeypointsOnImage.from_xy_array(
batch_pts[i], shape=batch_img[i].shape)
keypoints_after = KeypointsOnImage.from_xy_array(
aug_b_pts[i], shape=aug_b_imgs[i].shape)
image_before = keypoints_before.draw_on_image(batch_img[i])
image_after = keypoints_after.draw_on_image(aug_b_imgs[i])
ia.imshow(np.hstack([image_before, image_after]))
return aug_b_imgs, aug_b_pts
def __getitem__(self, item):
dp = self.dataset[item]
corners = dp['corners']
neighs = dp['neighs']
center = corners.mean(axis=0)
ng = (neighs - center) * 3 + center
mins = ng.min(axis=0).round().astype(int)
maxs = ng.max(axis=0).round().astype(int)
patch = dp['img'][mins[1]:maxs[1], mins[0]:maxs[0], :]
corners = corners - mins
neighs = neighs - mins
kps_c = KeypointsOnImage.from_xy_array(corners, patch.shape)
kps_n = KeypointsOnImage.from_xy_array(neighs, patch.shape)
kps = KeypointsOnImage(kps_c.keypoints + kps_n.keypoints, patch.shape)
if self.augment:
augs = [
iaa.Affine(rotate=(-10, 10)),
iaa.CropAndPad(percent=0.1),
iaa.Resize(64),
iaa.AdditiveGaussianNoise(scale=(0, 0.05 * 255))
]
else:
augs = [iaa.Resize(64)]
seq = iaa.Sequential(augs)
norm_patch, norm_kps = seq(image=patch, keypoints=kps)
norm_coords = norm_kps.to_xy_array().reshape([2, 4,
2]).transpose(0, 2, 1)
for dim in range(2):
norm_coords[:, dim, :] /= norm_patch.shape[dim]
#keypoints = np.concatenate([np.array([1, dp['color'], 1,1,1,1]), corners[:,0], corners[:,1], neighs[:,0], neighs[:,1]])
norm_keypoints = np.concatenate(
[np.array([1, dp['color'], 1, 1, 1, 1]),
norm_coords.reshape(-1)])
return torch.tensor(norm_patch).permute(
[2, 0, 1]).float() / 255, torch.tensor(norm_keypoints).float()
def test_augment_keypoints__kernel_size_is_two__no_keep_size(self):
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
kps = [Keypoint(x=1.5, y=5.5), Keypoint(x=5.5, y=1.5)]
kpsoi = KeypointsOnImage(kps, shape=(6, 6, 3))
expected = KeypointsOnImage.from_xy_array(np.float32(
[[1.5 / 2, 5.5 / 2], [5.5 / 2, 1.5 / 2]]),
shape=(3, 3, 3))
self._test_augment_cbaoi__kernel_size_is_two__no_keep_size(
kpsoi, expected, "augment_keypoints")
def test_augment_keypoints__kernel_size_is_two__no_keep_size(self):
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
kps = [Keypoint(x=1.5, y=5.5), Keypoint(x=5.5, y=1.5)]
kpsoi = KeypointsOnImage(kps, shape=(6, 6, 3))
aug = self.augmenter(2, keep_size=False)
kpsoi_aug = aug.augment_keypoints(kpsoi)
expected = KeypointsOnImage.from_xy_array(np.float32(
[[1.5 / 2, 5.5 / 2], [5.5 / 2, 1.5 / 2]]),
shape=(3, 3, 3))
assert_cbaois_equal(kpsoi_aug, expected)
def _test_augment_keypoints__kernel_size_differs(self, shape, shape_exp):
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
kps = [Keypoint(x=1.5, y=5.5), Keypoint(x=5.5, y=1.5)]
kpsoi = KeypointsOnImage(kps, shape=shape)
aug = self.augmenter((iap.Deterministic(3), iap.Deterministic(2)),
keep_size=False)
kpsoi_aug = aug.augment_keypoints(kpsoi)
expected = KeypointsOnImage.from_xy_array(np.float32([[
(1.5 / shape[1]) * shape_exp[1], (5.5 / shape[0]) * shape_exp[0]
], [(5.5 / shape[1]) * shape_exp[1],
(1.5 / shape[0]) * shape_exp[0]]]),
shape=shape_exp)
assert_cbaois_equal(kpsoi_aug, expected)
def augment_card(self, card, hull):
cardH, cardW, _ = card.shape
decalX = int((self.imgW - cardW) / 2)
decalY = int((self.imgH - cardH) / 2)
img_card = np.zeros((self.imgH, self.imgW, 4), dtype=np.uint8)
img_card[decalY:decalY + cardH, decalX:decalX + cardW, :] = card
card = img_card[:, :, :3]
hull = hull[:, 0, :]
card_kps_xy = np.float32([[decalX, decalY], [decalX + cardW, decalY],
[decalX + cardW, decalY + cardH],
[decalX, decalY + cardH]])
kpsoi_card = KeypointsOnImage.from_xy_array(card_kps_xy,
shape=card.shape)
# hull is a cv2.Contour, shape : Nx1x2
kpsoi_hull = [
ia.Keypoint(x=p[0] + decalX, y=p[1] + decalY)
for p in hull.reshape(-1, 2)
]
kpsoi_hull = KeypointsOnImage(kpsoi_hull,
shape=(self.imgH, self.imgW, 3))
# create polygon
poly_hull = Polygon(kpsoi_hull.keypoints)
# create empty segmentation map for classes: background and card
segmap = np.zeros((card.shape[0], card.shape[1], 3), dtype=np.uint8)
# draw the tree polygon into the second channel
segmap = poly_hull.draw_on_image(segmap,
color=(0, 255, 0),
alpha=1.0,
alpha_lines=0.0,
alpha_points=0.0)
# merge the two channels to a single one
segmap = np.argmax(segmap, axis=2)
segmap = segmap.astype(np.uint8)
segmap = SegmentationMapOnImage(segmap, nb_classes=2, shape=card.shape)
myseq = self.seq.to_deterministic()
card_aug, segmap_aug = myseq(image=card, segmentation_maps=segmap)
card_aug, kpsoi_aug = myseq(image=card, keypoints=kpsoi_card)
return card_aug, kpsoi_aug, segmap_aug
def test_augment_keypoints(self):
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
def do_return_augmentables(keypoints_on_images, parents, hooks):
return keypoints_on_images
aug_mock = mock.MagicMock(spec=meta.Augmenter)
aug_mock.augment_keypoints.side_effect = do_return_augmentables
kpsoi = KeypointsOnImage.from_xy_array(np.float32([[0, 0], [5, 1]]),
shape=(16, 24, 3))
aug = iaa.WithHueAndSaturation(aug_mock)
kpsoi_aug = aug.augment_keypoints(kpsoi)
assert kpsoi_aug.shape == (16, 24, 3)
assert kpsoi.keypoints[0].x == 0
assert kpsoi.keypoints[0].y == 0
assert kpsoi.keypoints[1].x == 5
assert kpsoi.keypoints[1].y == 1
assert aug_mock.augment_keypoints.call_count == 1
def normalize_keypoints(inputs, shapes=None):
# TODO get rid of this deferred import
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
shapes = _preprocess_shapes(shapes)
ntype = estimate_keypoints_norm_type(inputs)
_assert_exactly_n_shapes_partial = functools.partial(
_assert_exactly_n_shapes,
from_ntype=ntype,
to_ntype="List[KeypointsOnImage]",
shapes=shapes)
if ntype == "None":
return inputs
elif ntype in ["array[float]", "array[int]", "array[uint]"]:
_assert_single_array_ndim(inputs, 3, "(N,K,2)", "KeypointsOnImage")
_assert_single_array_last_dim_exactly(inputs, 2, "KeypointsOnImage")
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage.from_xy_array(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
elif ntype == "tuple[number,size=2]":
_assert_exactly_n_shapes_partial(n=1)
return [
KeypointsOnImage([Keypoint(x=inputs[0], y=inputs[1])],
shape=shapes[0])
]
elif ntype == "Keypoint":
_assert_exactly_n_shapes_partial(n=1)
return [KeypointsOnImage([inputs], shape=shapes[0])]
elif ntype == "KeypointsOnImage":
return [inputs]
elif ntype == "iterable[empty]":
return None
elif ntype in [
"iterable-array[float]", "iterable-array[int]",
"iterable-array[uint]"
]:
_assert_many_arrays_ndim(inputs, 2, "(K,2)", "KeypointsOnImage")
_assert_many_arrays_last_dim_exactly(inputs, 2, "KeypointsOnImage")
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage.from_xy_array(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
elif ntype == "iterable-tuple[number,size=2]":
_assert_exactly_n_shapes_partial(n=1)
return [
KeypointsOnImage([Keypoint(x=x, y=y) for x, y in inputs],
shape=shapes[0])
]
elif ntype == "iterable-Keypoint":
_assert_exactly_n_shapes_partial(n=1)
return [KeypointsOnImage(inputs, shape=shapes[0])]
elif ntype == "iterable-KeypointsOnImage":
return inputs
elif ntype == "iterable-iterable[empty]":
return None
elif ntype == "iterable-iterable-tuple[number,size=2]":
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage.from_xy_array(np.array(attr_i, dtype=np.float32),
shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
else:
assert ntype == "iterable-iterable-Keypoint", (
"Got unknown normalization type '%s'." % (ntype, ))
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage(attr_i, shape=shape)
for attr_i, shape in zip(inputs, shapes)
]
def normalize_keypoints(inputs, shapes=None):
# TODO get rid of this deferred import
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
shapes = _preprocess_shapes(shapes)
ntype = estimate_keypoints_norm_type(inputs)
_assert_exactly_n_shapes_partial = functools.partial(
_assert_exactly_n_shapes,
from_ntype=ntype, to_ntype="List[KeypointsOnImage]",
shapes=shapes)
if ntype == "None":
return inputs
elif ntype in ["array[float]", "array[int]", "array[uint]"]:
_assert_single_array_ndim(inputs, 3, "(N,K,2)", "KeypointsOnImage")
_assert_single_array_last_dim_exactly(inputs, 2, "KeypointsOnImage")
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage.from_xy_array(attr_i, shape=shape)
for attr_i, shape
in zip(inputs, shapes)
]
elif ntype == "tuple[number,size=2]":
_assert_exactly_n_shapes_partial(n=1)
return [KeypointsOnImage([Keypoint(x=inputs[0], y=inputs[1])],
shape=shapes[0])]
elif ntype == "Keypoint":
_assert_exactly_n_shapes_partial(n=1)
return [KeypointsOnImage([inputs], shape=shapes[0])]
elif ntype == "KeypointsOnImage":
return [inputs]
elif ntype == "iterable[empty]":
return None
elif ntype in ["iterable-array[float]",
"iterable-array[int]",
"iterable-array[uint]"]:
_assert_many_arrays_ndim(inputs, 2, "(K,2)", "KeypointsOnImage")
_assert_many_arrays_last_dim_exactly(inputs, 2, "KeypointsOnImage")
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage.from_xy_array(attr_i, shape=shape)
for attr_i, shape
in zip(inputs, shapes)
]
elif ntype == "iterable-tuple[number,size=2]":
_assert_exactly_n_shapes_partial(n=1)
return [KeypointsOnImage([Keypoint(x=x, y=y) for x, y in inputs],
shape=shapes[0])]
elif ntype == "iterable-Keypoint":
_assert_exactly_n_shapes_partial(n=1)
return [KeypointsOnImage(inputs, shape=shapes[0])]
elif ntype == "iterable-KeypointsOnImage":
return inputs
elif ntype == "iterable-iterable[empty]":
return None
elif ntype == "iterable-iterable-tuple[number,size=2]":
_assert_exactly_n_shapes_partial(n=len(inputs))
return [
KeypointsOnImage.from_xy_array(
np.array(attr_i, dtype=np.float32),
shape=shape)
for attr_i, shape
in zip(inputs, shapes)
]
else:
assert ntype == "iterable-iterable-Keypoint", (
"Got unknown normalization type '%s'." % (ntype,))
_assert_exactly_n_shapes_partial(n=len(inputs))
return [KeypointsOnImage(attr_i, shape=shape)
for attr_i, shape
in zip(inputs, shapes)]
kps = np.random.random((n_keypoints, 2))
kps[:, 0] *= image.shape[0]
kps[:, 1] *= image.shape[1]
keypoints.append(kps)
seq = iaa.Sequential(
[iaa.GaussianBlur((0, 3.0)),
iaa.Affine(scale=(0.5, 0.7))])
# augment keypoints and images
images_aug, keypoints_aug = seq(images=images, keypoints=keypoints)
# Example code to show each image and print the new keypoints coordinates
for i in range(len(images)):
print("[Image #%d]" % (i, ))
keypoints_before = KeypointsOnImage.from_xy_array(keypoints[i],
shape=images[i].shape)
keypoints_after = KeypointsOnImage.from_xy_array(keypoints_aug[i],
shape=images_aug[i].shape)
image_before = keypoints_before.draw_on_image(images[i])
image_after = keypoints_after.draw_on_image(images_aug[i])
ia.imshow(np.hstack([image_before, image_after]))
kps_zipped = zip(keypoints_before.keypoints, keypoints_after.keypoints)
for keypoint_before, keypoint_after in kps_zipped:
x_before, y_before = keypoint_before.x, keypoint_before.y
x_after, y_after = keypoint_after.x, keypoint_after.y
print("before aug: x=%d y=%d | after aug: x=%d y=%d" %
(x_before, y_before, x_after, y_after))
def generator(self, features, labels, batch_size):
batch_features = np.zeros((batch_size, 128, 128, 3))
