mirror of
https://github.com/open-mmlab/mmocr.git
synced 2025-06-03 21:54:47 +08:00
47896a3f80
* Add unit test for dbnet transforms * Add unit test for dbnet transforms * Add unit test for dbnet transforms
36 lines
1.4 KiB
Python
36 lines
1.4 KiB
Python
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())
|