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}