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