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 kps_to_mask(self, kpsoi):
poly_card = Polygon(kpsoi.keypoints)
segmap = np.zeros((self.imgH, self.imgW, 3), dtype=np.uint8)
# draw the tree polygon into the second channel
segmap = poly_card.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)
card_mask = np.stack([segmap] * 3, -1)
return card_mask
def __getitem__(self, idx):
image = Image.open(os.path.join(
self.root, self.image_filenames[idx])).convert('RGB')
image = np.asarray(image)
if self.nums is not None:
segmap = np.zeros((image.shape[0], image.shape[1], 1),
dtype=np.int32)
nums = self.nums[idx]
for num in nums:
poly = Polygon(num['box'])
segmap = poly.draw_on_image(segmap, color=2)
segmap = SegmentationMapsOnImage(segmap, shape=image.shape)
if self.transforms:
image, segmap = self.transforms(image=image,
segmentation_maps=segmap)
image = to_tensor(image)
image = normalize(image, MEAN, STD)
segmap = segmap.arr
segmap = to_tensor(segmap.astype(np.float32))
return {'image': image, 'mask': segmap}
else:
if self.transforms:
image = self.transforms(image=image)
image = to_tensor(image)
image = normalize(image, MEAN, STD)
return {'image': image}