mirror of https://github.com/open-mmlab/mmocr.git
283 lines
12 KiB
Python
283 lines
12 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import copy
|
|
import unittest
|
|
import unittest.mock as mock
|
|
|
|
import numpy as np
|
|
|
|
from mmocr.datasets.pipelines import (PyramidRescale, RandomRotate, Resize,
|
|
TextDetRandomCropFlip)
|
|
|
|
|
|
class TestPyramidRescale(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.data_info = dict(img=np.random.random((128, 100, 3)))
|
|
|
|
def test_init(self):
|
|
# factor is int
|
|
transform = PyramidRescale(factor=4, randomize_factor=False)
|
|
self.assertEqual(transform.factor, 4)
|
|
# factor is float
|
|
with self.assertRaisesRegex(TypeError,
|
|
'`factor` should be an integer'):
|
|
PyramidRescale(factor=4.0)
|
|
# invalid base_shape
|
|
with self.assertRaisesRegex(TypeError,
|
|
'`base_shape` should be a list or tuple'):
|
|
PyramidRescale(base_shape=128)
|
|
with self.assertRaisesRegex(
|
|
ValueError, '`base_shape` should contain two integers'):
|
|
PyramidRescale(base_shape=(128, ))
|
|
with self.assertRaisesRegex(
|
|
ValueError, '`base_shape` should contain two integers'):
|
|
PyramidRescale(base_shape=(128.0, 2.0))
|
|
# invalid randomize_factor
|
|
with self.assertRaisesRegex(TypeError,
|
|
'`randomize_factor` should be a bool'):
|
|
PyramidRescale(randomize_factor=None)
|
|
|
|
def test_transform(self):
|
|
# test if the rescale keeps the original size
|
|
transform = PyramidRescale()
|
|
results = transform(copy.deepcopy(self.data_info))
|
|
self.assertEqual(results['img'].shape, (128, 100, 3))
|
|
# test factor = 0
|
|
transform = PyramidRescale(factor=0, randomize_factor=False)
|
|
results = transform(copy.deepcopy(self.data_info))
|
|
self.assertTrue(np.all(results['img'] == self.data_info['img']))
|
|
|
|
def test_repr(self):
|
|
transform = PyramidRescale(
|
|
factor=4, base_shape=(128, 512), randomize_factor=False)
|
|
self.assertEqual(
|
|
repr(transform),
|
|
('PyramidRescale(factor = 4, randomize_factor = False, '
|
|
'base_w = 128, base_h = 512)'))
|
|
|
|
|
|
class TestTextDetRandomCropFlip(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
img = np.ones((10, 10, 3))
|
|
img[0, 0, :] = 0
|
|
self.data_info1 = dict(
|
|
img=copy.deepcopy(img),
|
|
gt_polygons=[np.array([0., 0., 0., 10., 10., 10., 10., 0.])],
|
|
img_shape=[10, 10])
|
|
self.data_info2 = dict(
|
|
img=copy.deepcopy(img),
|
|
gt_polygons=[np.array([1., 1., 1., 9., 9., 9., 9., 1.])],
|
|
img_shape=[10, 10])
|
|
|
|
def test_init(self):
|
|
# iter_num is int
|
|
transform = TextDetRandomCropFlip(iter_num=1)
|
|
self.assertEqual(transform.iter_num, 1)
|
|
# iter_num is float
|
|
with self.assertRaisesRegex(TypeError,
|
|
'`iter_num` should be an integer'):
|
|
transform = TextDetRandomCropFlip(iter_num=1.5)
|
|
|
|
@mock.patch('mmocr.datasets.pipelines.processing.np.random.randint')
|
|
def test_transforms(self, mock_sample):
|
|
mock_sample.side_effect = [0, 1, 2]
|
|
transform = TextDetRandomCropFlip(crop_ratio=1.0, iter_num=3)
|
|
results = transform(self.data_info2)
|
|
self.assertTrue(np.allclose(results['img'], self.data_info2['img']))
|
|
self.assertTrue(
|
|
np.allclose(results['gt_polygons'],
|
|
self.data_info2['gt_polygons']))
|
|
|
|
def test_generate_crop_target(self):
|
|
transform = TextDetRandomCropFlip(
|
|
crop_ratio=1.0, iter_num=3, pad_ratio=0.1)
|
|
h, w = self.data_info1['img_shape']
|
|
pad_h = int(h * transform.pad_ratio)
|
|
pad_w = int(w * transform.pad_ratio)
|
|
h_axis, w_axis = transform._generate_crop_target(
|
|
self.data_info1['img'], self.data_info1['gt_polygons'], pad_h,
|
|
pad_w)
|
|
self.assertTrue(np.allclose(h_axis, (0, 11)))
|
|
self.assertTrue(np.allclose(w_axis, (0, 11)))
|
|
|
|
def test_repr(self):
|
|
transform = TextDetRandomCropFlip(
|
|
pad_ratio=0.1,
|
|
crop_ratio=0.5,
|
|
iter_num=1,
|
|
min_area_ratio=0.2,
|
|
epsilon=1e-2)
|
|
self.assertEqual(
|
|
repr(transform),
|
|
('TextDetRandomCropFlip(pad_ratio = 0.1, crop_ratio = 0.5, '
|
|
'iter_num = 1, min_area_ratio = 0.2, epsilon = 0.01)'))
|
|
|
|
|
|
class TestRandomRotate(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
img = np.random.random((5, 5))
|
|
self.data_info1 = dict(img=img.copy())
|
|
self.data_info2 = dict(
|
|
img=np.random.random((30, 30, 3)),
|
|
gt_bboxes=np.array([[10, 10, 20, 20], [5, 5, 10, 10]]))
|
|
self.data_info3 = dict(
|
|
img=np.random.random((30, 30, 3)),
|
|
gt_polygons=[np.array([10., 10., 20., 10., 20., 20., 10., 20.])])
|
|
|
|
def test_init(self):
|
|
# max angle is float
|
|
with self.assertRaisesRegex(TypeError,
|
|
'`max_angle` should be an integer'):
|
|
RandomRotate(max_angle=16.8)
|
|
# invalid pad value
|
|
with self.assertRaisesRegex(
|
|
ValueError, '`pad_value` should contain three integers'):
|
|
RandomRotate(pad_value=[16.8, 0.1])
|
|
|
|
def test_transform(self):
|
|
self._test_recog()
|
|
self._test_bboxes()
|
|
self._test_polygons()
|
|
|
|
def _test_recog(self):
|
|
# test random rotate for recognition (image only) input
|
|
transform = RandomRotate(max_angle=10)
|
|
results = transform(copy.deepcopy(self.data_info1))
|
|
self.assertTrue(np.allclose(results['img'], self.data_info1['img']))
|
|
|
|
@mock.patch('mmocr.datasets.pipelines.processing.np.random.random_sample')
|
|
def _test_bboxes(self, mock_sample):
|
|
# test random rotate for bboxes
|
|
# returns 1. for random_sample() in _sample_angle(), i.e., angle = 90
|
|
mock_sample.side_effect = [1.]
|
|
transform = RandomRotate(max_angle=90, use_canvas=True)
|
|
results = transform(copy.deepcopy(self.data_info2))
|
|
self.assertTrue(
|
|
np.allclose(results['gt_bboxes'][0], np.array([10, 10, 20, 20])))
|
|
self.assertTrue(
|
|
np.allclose(results['gt_bboxes'][1], np.array([5, 20, 10, 25])))
|
|
self.assertEqual(results['img'].shape, self.data_info2['img'].shape)
|
|
|
|
@mock.patch('mmocr.datasets.pipelines.processing.np.random.random_sample')
|
|
def _test_polygons(self, mock_sample):
|
|
# test random rotate for polygons
|
|
# returns 1. for random_sample() in _sample_angle(), i.e., angle = 90
|
|
mock_sample.side_effect = [1.]
|
|
transform = RandomRotate(max_angle=90, use_canvas=True)
|
|
results = transform(copy.deepcopy(self.data_info3))
|
|
self.assertTrue(
|
|
np.allclose(results['gt_polygons'][0],
|
|
np.array([10., 20., 10., 10., 20., 10., 20., 20.])))
|
|
self.assertEqual(results['img'].shape, self.data_info3['img'].shape)
|
|
|
|
def test_repr(self):
|
|
transform = RandomRotate(
|
|
max_angle=10,
|
|
pad_with_fixed_color=False,
|
|
pad_value=(0, 0, 0),
|
|
use_canvas=False)
|
|
self.assertEqual(
|
|
repr(transform),
|
|
('RandomRotate(max_angle = 10, '
|
|
'pad_with_fixed_color = False, pad_value = (0, 0, 0), '
|
|
'use_canvas = False)'))
|
|
|
|
|
|
class TestResize(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.data_info1 = dict(
|
|
img=np.random.random((600, 800, 3)),
|
|
gt_bboxes=np.array([[0, 0, 60, 100]]),
|
|
gt_polygons=[np.array([0, 0, 200, 0, 200, 100, 0, 100])])
|
|
self.data_info2 = dict(
|
|
img=np.random.random((200, 300, 3)),
|
|
gt_bboxes=np.array([[0, 0, 400, 600]]),
|
|
gt_polygons=[np.array([0, 0, 400, 0, 400, 400, 0, 400])])
|
|
self.data_info3 = dict(
|
|
img=np.random.random((200, 300, 3)),
|
|
gt_bboxes=np.array([[400, 400, 600, 600]]),
|
|
gt_polygons=[np.array([400, 400, 500, 400, 500, 600, 400, 600])])
|
|
|
|
def test_resize(self):
|
|
# test keep_ratio is True
|
|
transform = Resize(scale=(400, 400), keep_ratio=True)
|
|
results = transform(copy.deepcopy(self.data_info1.copy()))
|
|
self.assertEqual(results['img'].shape[:2], (300, 400))
|
|
self.assertEqual(results['img_shape'], (300, 400))
|
|
self.assertEqual(results['scale'], (400, 300))
|
|
self.assertEqual(results['scale_factor'], (400 / 800, 300 / 600))
|
|
self.assertEqual(results['keep_ratio'], True)
|
|
|
|
# test keep_ratio is False
|
|
transform = Resize(scale=(400, 400))
|
|
results = transform(copy.deepcopy(self.data_info1.copy()))
|
|
self.assertEqual(results['img'].shape[:2], (400, 400))
|
|
self.assertEqual(results['img_shape'], (400, 400))
|
|
self.assertEqual(results['scale'], (400, 400))
|
|
self.assertEqual(results['scale_factor'], (400 / 800, 400 / 600))
|
|
self.assertEqual(results['keep_ratio'], False)
|
|
|
|
# test resize_bboxes/polygons
|
|
transform = Resize(scale_factor=(1.5, 2))
|
|
results = transform(copy.deepcopy(self.data_info1.copy()))
|
|
self.assertEqual(results['img'].shape[:2], (1200, 1200))
|
|
self.assertEqual(results['img_shape'], (1200, 1200))
|
|
self.assertEqual(results['scale'], (1200, 1200))
|
|
self.assertEqual(results['scale_factor'], (1.5, 2))
|
|
self.assertEqual(results['keep_ratio'], False)
|
|
self.assertTrue(
|
|
results['gt_bboxes'].all() == np.array([[0, 0, 90, 200]]).all())
|
|
self.assertTrue(results['gt_polygons'][0].all() == np.array(
|
|
[0, 0, 300, 0, 300, 200, 0, 200]).all())
|
|
|
|
# test clip_object_border = False
|
|
transform = Resize(scale=(150, 100), clip_object_border=False)
|
|
results = transform(self.data_info2.copy())
|
|
self.assertEqual(results['img'].shape[:2], (100, 150))
|
|
self.assertEqual(results['img_shape'], (100, 150))
|
|
self.assertEqual(results['scale'], (150, 100))
|
|
self.assertEqual(results['scale_factor'], (150. / 300., 100. / 200.))
|
|
self.assertEqual(results['keep_ratio'], False)
|
|
self.assertTrue(
|
|
results['gt_bboxes'].all() == np.array([0, 0, 200, 300]).all())
|
|
self.assertTrue(results['gt_polygons'][0].all() == np.array(
|
|
[0, 0, 200, 0, 200, 200, 0, 200]).all())
|
|
|
|
# test clip_object_border = True
|
|
transform = Resize(scale=(150, 100), clip_object_border=True)
|
|
results = transform(self.data_info2.copy())
|
|
self.assertEqual(results['img'].shape[:2], (100, 150))
|
|
self.assertEqual(results['img_shape'], (100, 150))
|
|
self.assertEqual(results['scale'], (150, 100))
|
|
self.assertEqual(results['scale_factor'], (150. / 300., 100. / 200.))
|
|
self.assertEqual(results['keep_ratio'], False)
|
|
self.assertTrue(
|
|
results['gt_bboxes'].all() == np.array([0, 0, 150, 100]).all())
|
|
self.assertTrue(results['gt_polygons'][0].shape == (8, ))
|
|
self.assertTrue(results['gt_polygons'][0].all() == np.array(
|
|
[0, 0, 150, 0, 150, 100, 0, 100]).all())
|
|
|
|
# test clip_object_border = True and polygon outside image
|
|
transform = Resize(scale=(150, 100), clip_object_border=True)
|
|
results = transform(self.data_info3)
|
|
self.assertEqual(results['img'].shape[:2], (100, 150))
|
|
self.assertEqual(results['img_shape'], (100, 150))
|
|
self.assertEqual(results['scale'], (150, 100))
|
|
self.assertEqual(results['scale_factor'], (150. / 300., 100. / 200.))
|
|
self.assertEqual(results['keep_ratio'], False)
|
|
self.assertEqual(results['gt_polygons'][0].all(),
|
|
np.array([0., 0., 0., 0., 0., 0., 0., 0.]).all())
|
|
self.assertEqual(results['gt_bboxes'].all(),
|
|
np.array([[150., 100., 150., 100.]]).all())
|
|
|
|
def test_repr(self):
|
|
transform = Resize(scale=(2000, 2000), keep_ratio=True)
|
|
self.assertEqual(
|
|
repr(transform), ('Resize(scale=(2000, 2000), '
|
|
'scale_factor=None, keep_ratio=True, '
|
|
'clip_object_border=True), backend=cv2), '
|
|
'interpolation=bilinear)'))
|