|
|
|
|
|
import numpy as np |
|
|
|
|
|
import mmocr.datasets.pipelines.dbnet_transforms as transforms |
|
|
|
|
|
|
|
|
def test_imgaug(): |
|
|
args = [['Fliplr', 0.5], |
|
|
dict(cls='Affine', rotate=[-10, 10]), ['Resize', [0.5, 3.0]]] |
|
|
imgaug = transforms.ImgAug(args) |
|
|
img = np.random.rand(3, 100, 200) |
|
|
poly = np.array([[[0, 0, 50, 0, 50, 50, 0, 50]], |
|
|
[[20, 20, 50, 20, 50, 50, 20, 50]]]) |
|
|
box = np.array([[0, 0, 50, 50], [20, 20, 50, 50]]) |
|
|
results = dict(img=img, masks=poly, bboxes=box) |
|
|
results['mask_fields'] = ['masks'] |
|
|
results['bbox_fields'] = ['bboxes'] |
|
|
results = imgaug(results) |
|
|
assert np.allclose(results['bboxes'][0], |
|
|
results['masks'].masks[0][0][[0, 1, 4, 5]]) |
|
|
assert np.allclose(results['bboxes'][1], |
|
|
results['masks'].masks[1][0][[0, 1, 4, 5]]) |
|
|
|
|
|
|
|
|
def test_eastrandomcrop(): |
|
|
crop = transforms.EastRandomCrop(target_size=(60, 60), max_tries=100) |
|
|
img = np.random.rand(3, 100, 200) |
|
|
poly = np.array([[[0, 0, 50, 0, 50, 50, 0, 50]], |
|
|
[[20, 20, 50, 20, 50, 50, 20, 50]]]) |
|
|
box = np.array([[0, 0, 50, 50], [20, 20, 50, 50]]) |
|
|
results = dict(img=img, gt_masks=poly, bboxes=box) |
|
|
results['mask_fields'] = ['gt_masks'] |
|
|
results['bbox_fields'] = ['bboxes'] |
|
|
results = crop(results) |
|
|
assert np.allclose(results['bboxes'][0], |
|
|
results['gt_masks'].masks[0][0][[0, 2]].flatten()) |
|
|
|
