def test_transform(): data_root = '/home/luning/dev/data/SynthText800k/detection/' img_fn = 'desert_78_83' gt_fn = data_root + 'gt.mat' targets = {} targets = sio.loadmat(gt_fn, targets, squeeze_me=True, struct_as_record=False, variable_names=['imnames', 'wordBB', 'txt']) imageNames = targets['imnames'] wordBBoxes = targets['wordBB'] transcripts = targets['txt'] mask = [True if img_fn in i else False for i in imageNames] index = np.where(mask)[0][0] print(index) img_fn = imageNames[index] img_fn = data_root + '/imgs/' + img_fn img = cv2.imread(img_fn) boxes = wordBBoxes[index] transcripts = transcripts[index] transcripts = [word for line in transcripts for word in line.split()] boxes = np.expand_dims(boxes, axis=2) if (boxes.ndim == 2) else boxes _, _, numOfWords = boxes.shape # boxes = boxes.reshape([8, numOfWords]).T # num_words * 8 # boxes = boxes.reshape([numOfWords, 4, 2]) boxes = boxes.transpose((2, 1, 0)) polys = [] for i in range(numOfWords): # box = boxes[:, :, i].T box = boxes[i] polys.append(Polygon(box.tolist())) polys_on_image = PolygonsOnImage(polygons=polys, shape=img.shape) image_before_aug = polys_on_image.draw_on_image(img) cv2.imwrite('before_aug.jpg', image_before_aug) transform = Transform() image_aug, polygons_aug = transform(img, boxes) polygons_on_image = PolygonsOnImage(polygons=polygons_aug, shape=image_aug.shape) image_aug = polygons_on_image.draw_on_image(image_aug) cv2.imwrite('aug.jpg', image_aug)
def aug(self,im,bboxes,viz=True): assert len(bboxes.shape)==3 assert bboxes.shape[1]==4 poly_on_img = PolygonsOnImage([Polygon(bbox) for bbox in bboxes], shape=im.shape) if viz: res = poly_on_img.draw_on_image(im) cv2.imshow('ori', res) cv2.waitKey(0) imgs_aug, poly_on_img_aug = self.seq(image=im, polygons=poly_on_img,) res_bboxes=poly_on_img_aug.to_xy_array() if bboxes.shape[0]*bboxes.shape[1]!=res_bboxes.shape[0]: print('aug error','before:',bboxes.shape,' after:',res_bboxes.shape) return im,bboxes #print('before:',bboxes.shape,' after:',res_bboxes.shape) res_bboxes=np.reshape(res_bboxes,bboxes.shape) if viz: res = poly_on_img_aug.draw_on_image(imgs_aug) cv2.imshow('a', res) cv2.waitKey(0) return imgs_aug,res_bboxes