# Copyright (c) OpenMMLab. All rights reserved. import copy import unittest import unittest.mock as mock import numpy as np from mmcv.transforms import Pad, RandomResize from mmocr.datasets.pipelines import (BoundedScaleAspectJitter, PadToWidth, PyramidRescale, RandomCrop, RandomFlip, RandomRotate, RescaleToHeight, Resize, ShortScaleAspectJitter, SourceImagePad, TextDetRandomCrop, TextDetRandomCropFlip) from mmocr.utils import bbox2poly, poly2shapely 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.])], gt_bboxes_labels=np.array([0], dtype=np.int64), gt_ignored=np.array([True], dtype=np.bool_), img_shape=[10, 10]) self.data_info3 = dict( img=copy.deepcopy(img), gt_polygons=[ np.array([0., 0., 4., 0., 4., 4., 0., 4.]), np.array([4., 0., 8., 0., 8., 4., 4., 4.]) ], gt_bboxes_labels=np.array([0, 0], dtype=np.int64), gt_ignored=np.array([True, True], dtype=np.bool_), 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'])) self.assertEqual( len(results['gt_bboxes']), len(results['gt_polygons'])) self.assertTrue( poly2shapely(results['gt_polygons'][0]).equals( poly2shapely(bbox2poly(results['gt_bboxes'][0])))) def test_size(self): transform = TextDetRandomCropFlip(crop_ratio=1.0, iter_num=3) results = transform(self.data_info3) self.assertEqual( len(results['gt_bboxes']), len(results['gt_polygons'])) self.assertEqual( len(results['gt_polygons']), len(results['gt_ignored'])) self.assertEqual( len(results['gt_ignored']), len(results['gt_bboxes_labels'])) 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(), img_shape=img.shape[:2]) self.data_info2 = dict( img=np.random.random((30, 30, 3)), gt_bboxes=np.array([[10, 10, 20, 20], [5, 5, 10, 10]]), img_shape=(30, 30)) self.data_info3 = dict( img=np.random.random((30, 30, 3)), gt_polygons=[np.array([10., 10., 20., 10., 20., 20., 10., 20.])], img_shape=(30, 30)) 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 TestRandomCrop(unittest.TestCase): def setUp(self): img = np.zeros((30, 30, 3)) gt_polygons = [ np.array([5., 5., 25., 5., 25., 10., 5., 10.]), np.array([5., 20., 25., 20., 25., 25., 5., 25.]) ] gt_bboxes = np.array([[5, 5, 25, 10], [5, 20, 25, 25]]) labels = np.array([0, 1]) gt_ignored = np.array([True, False], dtype=bool) texts = ['text1', 'text2'] self.data_info = dict( img=img, gt_polygons=gt_polygons, gt_bboxes=gt_bboxes, gt_bboxes_labels=labels, gt_ignored=gt_ignored, gt_texts=texts) @mock.patch('mmocr.datasets.pipelines.processing.np.random.randint') def test_sample_crop_box(self, mock_randint): def rand_min(low, high): return low trans = RandomCrop(min_side_ratio=0.3) mock_randint.side_effect = rand_min crop_box = trans._sample_crop_box((30, 30), self.data_info.copy()) assert np.allclose(np.array(crop_box), np.array([0, 0, 25, 10])) def rand_max(low, high): return high - 1 mock_randint.side_effect = rand_max crop_box = trans._sample_crop_box((30, 30), self.data_info.copy()) assert np.allclose(np.array(crop_box), np.array([4, 19, 30, 30])) @mock.patch('mmocr.datasets.pipelines.processing.np.random.randint') def test_transform(self, mock_randint): def rand_min(low, high): return low # mock_randint.side_effect = [0, 0, 0, 0, 30, 0, 0, 0, 15] mock_randint.side_effect = rand_min trans = RandomCrop(min_side_ratio=0.3) polygon_target = np.array([5., 5., 25., 5., 25., 10., 5., 10.]) bbox_target = np.array([[5., 5., 25., 10.]]) results = trans(self.data_info) self.assertEqual(results['img'].shape, (10, 25, 3)) self.assertEqual(results['img_shape'], (10, 25)) self.assertTrue(np.allclose(results['gt_bboxes'], bbox_target)) self.assertEqual(results['gt_bboxes'].shape, (1, 4)) self.assertEqual(len(results['gt_polygons']), 1) self.assertTrue(np.allclose(results['gt_polygons'][0], polygon_target)) self.assertEqual(results['gt_bboxes_labels'][0], 0) self.assertEqual(results['gt_ignored'][0], True) self.assertEqual(results['gt_texts'][0], 'text1') def rand_max(low, high): return high - 1 mock_randint.side_effect = rand_max trans = RandomCrop(min_side_ratio=0.3) polygon_target = np.array([1, 1, 21, 1, 21, 6, 1, 6]) bbox_target = np.array([[1, 1, 21, 6]]) results = trans(self.data_info) self.assertEqual(results['img'].shape, (6, 21, 3)) self.assertEqual(results['img_shape'], (6, 21)) self.assertTrue(np.allclose(results['gt_bboxes'], bbox_target)) self.assertEqual(results['gt_bboxes'].shape, (1, 4)) self.assertEqual(len(results['gt_polygons']), 1) self.assertTrue(np.allclose(results['gt_polygons'][0], polygon_target)) self.assertEqual(results['gt_bboxes_labels'][0], 0) self.assertTrue(results['gt_ignored'][0]) self.assertEqual(results['gt_texts'][0], 'text1') def test_repr(self): transform = RandomCrop(min_side_ratio=0.4) self.assertEqual(repr(transform), ('RandomCrop(min_side_ratio = 0.4)')) class TestTextDetRandomCrop(unittest.TestCase): def setUp(self): img = np.array([[[1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]]).transpose(1, 2, 0) gt_polygons = [np.array([2, 2, 5, 2, 5, 5, 2, 5])] gt_bboxes = np.array([[2, 2, 5, 5]]) gt_bboxes_labels = np.array([0]) gt_ignored = np.array([True], dtype=bool) self.data_info = dict( img=img, gt_polygons=gt_polygons, gt_bboxes=gt_bboxes, gt_bboxes_labels=gt_bboxes_labels, gt_ignored=gt_ignored) @mock.patch('mmocr.datasets.pipelines.processing.np.random.random_sample') @mock.patch('mmocr.datasets.pipelines.processing.random.randint') def test_sample_offset(self, mock_randint, mock_sample): # test target size is bigger than image size mock_sample.side_effect = [1] trans = TextDetRandomCrop(target_size=(6, 6)) offset = trans._sample_offset(self.data_info['gt_polygons'], self.data_info['img'].shape[:2]) self.assertEqual(offset, (0, 0)) # test the first bracnh in sample_offset mock_sample.side_effect = [0.1] mock_randint.side_effect = [0, 2] trans = TextDetRandomCrop(target_size=(3, 3)) offset = trans._sample_offset(self.data_info['gt_polygons'], self.data_info['img'].shape[:2]) self.assertEqual(offset, (0, 2)) # test the second branch in sample_offset mock_sample.side_effect = [1] mock_randint.side_effect = [1, 2] trans = TextDetRandomCrop(target_size=(3, 3)) offset = trans._sample_offset(self.data_info['gt_polygons'], self.data_info['img'].shape[:2]) self.assertEqual(offset, (1, 2)) mock_sample.side_effect = [1] mock_randint.side_effect = [1, 2] trans = TextDetRandomCrop(target_size=(5, 5)) offset = trans._sample_offset(self.data_info['gt_polygons'], self.data_info['img'].shape[:2]) self.assertEqual(offset, (0, 0)) def test_crop_image(self): img = self.data_info['img'] offset = [0, 0] target = [6, 6] trans = TextDetRandomCrop(target_size=(3, 3)) crop, _ = trans._crop_img(img, offset, target) self.assertEqual(img.shape, crop.shape) target = [3, 2] crop = trans._crop_img(img, offset, target) self.assertTrue( np.allclose( np.array([[[1, 2, 3], [1, 2, 3]]]).transpose(1, 2, 0), crop[0])) self.assertTrue(np.allclose(crop[1], np.array([0, 0, 3, 2]))) def test_crop_polygons(self): trans = TextDetRandomCrop(target_size=(3, 3)) crop_box = np.array([2, 3, 5, 5]) polygons = [ bbox2poly([2, 3, 4, 4]), bbox2poly([0, 0, 1, 1]), bbox2poly([1, 2, 4, 4]), bbox2poly([0, 0, 10, 10]) ] kept_polygons, kept_idx = trans._crop_polygons(polygons, crop_box) target_polygons = [ bbox2poly([0, 0, 2, 1]), bbox2poly([0, 0, 2, 1]), bbox2poly([0, 0, 3, 2]), ] self.assertEqual(len(kept_polygons), 3) self.assertEqual(kept_idx, [0, 2, 3]) self.assertTrue( poly2shapely(target_polygons[0]).equals( poly2shapely(kept_polygons[0]))) self.assertTrue( poly2shapely(target_polygons[1]).equals( poly2shapely(kept_polygons[1]))) self.assertTrue( poly2shapely(target_polygons[2]).equals( poly2shapely(kept_polygons[2]))) @mock.patch('mmocr.datasets.pipelines.processing.np.random.random_sample') @mock.patch('mmocr.datasets.pipelines.processing.random.randint') def test_transform(self, mock_randint, mock_sample): # test target size is equal to image size trans = TextDetRandomCrop(target_size=(5, 5)) data_info = self.data_info.copy() results = trans(data_info) self.assertDictEqual(results, data_info) mock_sample.side_effect = [0.1] mock_randint.side_effect = [1, 1] trans = TextDetRandomCrop(target_size=(3, 3)) data_info = self.data_info.copy() results = trans(data_info) box_target = np.array([1, 1, 3, 3]) polygon_target = np.array([1, 1, 3, 1, 3, 3, 1, 3]) self.assertEqual(results['img'].shape, (3, 3, 1)) self.assertEqual(results['img_shape'], (3, 3)) self.assertTrue( poly2shapely(bbox2poly(box_target)).equals( poly2shapely(bbox2poly(results['gt_bboxes'][0])))) self.assertTrue( poly2shapely(polygon_target).equals( poly2shapely(results['gt_polygons'][0]))) self.assertTrue(results['gt_bboxes_labels'] == np.array([0])) self.assertTrue(results['gt_ignored'][0]) def test_repr(self): transform = TextDetRandomCrop( target_size=(512, 512), positive_sample_ratio=0.4) self.assertEqual( repr(transform), ('TextDetRandomCrop(target_size = (512, 512), ' 'positive_sample_ratio = 0.4)')) class TestEastRandomCrop(unittest.TestCase): def setUp(self): img = np.ones((30, 30, 3)) gt_polygons = [ np.array([5., 5., 25., 5., 25., 10., 5., 10.]), np.array([5., 20., 25., 20., 25., 25., 5., 25.]) ] gt_bboxes = np.array([[5, 5, 25, 10], [5, 20, 25, 25]]) labels = np.array([0, 1]) gt_ignored = np.array([True, False], dtype=bool) texts = ['text1', 'text2'] self.data_info = dict( img=img, gt_polygons=gt_polygons, gt_bboxes=gt_bboxes, gt_bboxes_labels=labels, gt_ignored=gt_ignored, gt_texts=texts) @mock.patch('mmocr.datasets.pipelines.processing.np.random.randint') def test_east_random_crop(self, mock_randint): # test randomcrop randcrop = RandomCrop(min_side_ratio=0.5) mock_randint.side_effect = [0, 0, 0, 0, 30, 0, 0, 0, 15] crop_results = randcrop(self.data_info) polygon_target = np.array([5., 5., 25., 5., 25., 10., 5., 10.]) bbox_target = np.array([[5., 5., 25., 10.]]) self.assertEqual(crop_results['img'].shape, (15, 30, 3)) self.assertEqual(crop_results['img_shape'], (15, 30)) self.assertTrue(np.allclose(crop_results['gt_bboxes'], bbox_target)) self.assertEqual(crop_results['gt_bboxes'].shape, (1, 4)) self.assertEqual(len(crop_results['gt_polygons']), 1) self.assertTrue( np.allclose(crop_results['gt_polygons'][0], polygon_target)) self.assertEqual(crop_results['gt_bboxes_labels'][0], 0) self.assertTrue(crop_results['gt_ignored'][0]) self.assertEqual(crop_results['gt_texts'][0], 'text1') # test resize resize = Resize(scale=(30, 30), keep_ratio=True) resize_results = resize(crop_results) self.assertEqual(resize_results['img'].shape, (15, 30, 3)) self.assertEqual(crop_results['img_shape'], (15, 30)) self.assertEqual(crop_results['scale'], (30, 30)) self.assertEqual(crop_results['scale_factor'], (1., 1.)) self.assertTrue(crop_results['keep_ratio']) # test pad pad = Pad(size=(30, 30)) pad_results = pad(resize_results) self.assertEqual(pad_results['img'].shape, (30, 30, 3)) self.assertEqual(pad_results['pad_shape'], (30, 30, 3)) self.assertEqual(pad_results['img'].sum(), 15 * 30 * 3) class TestRandomResize(unittest.TestCase): def setUp(self): self.data_info1 = dict( img=np.random.random((300, 400, 3)), gt_bboxes=np.array([[0, 0, 60, 100]]), gt_polygons=[np.array([0, 0, 200, 0, 200, 100, 0, 100])]) @mock.patch('mmcv.transforms.processing.np.random.random_sample') def test_random_resize(self, mock_sample): randresize = RandomResize( scale=(500, 500), ratio_range=(0.8, 1.2), resize_type='mmocr.Resize', keep_ratio=True) target_bboxes = np.array([0, 0, 90, 150]) target_polygons = [np.array([0, 0, 300, 0, 300, 150, 0, 150])] mock_sample.side_effect = [1.0] results = randresize(self.data_info1) self.assertEqual(results['img'].shape, (450, 600, 3)) self.assertEqual(results['img_shape'], (450, 600)) self.assertEqual(results['keep_ratio'], True) self.assertEqual(results['scale'], (600, 600)) self.assertEqual(results['scale_factor'], (600. / 400., 450. / 300.)) self.assertTrue( poly2shapely(bbox2poly(results['gt_bboxes'][0])).equals( poly2shapely(bbox2poly(target_bboxes)))) self.assertTrue( poly2shapely(results['gt_polygons'][0]).equals( poly2shapely(target_polygons[0]))) class TestRescaleToHeight(unittest.TestCase): def test_rescale_height(self): data_info = dict( img=np.random.random((16, 25, 3)), gt_seg_map=np.random.random((16, 25, 3)), gt_bboxes=np.array([[0, 0, 10, 10]]), gt_keypoints=np.array([[[10, 10, 1]]])) with self.assertRaises(AssertionError): RescaleToHeight(height=20.9) with self.assertRaises(AssertionError): RescaleToHeight(height=20, min_width=20.9) with self.assertRaises(AssertionError): RescaleToHeight(height=20, max_width=20.9) with self.assertRaises(AssertionError): RescaleToHeight(height=20, width_divisor=0.5) transform = RescaleToHeight(height=32) results = transform(copy.deepcopy(data_info)) self.assertTupleEqual(results['img'].shape[:2], (32, 50)) self.assertTupleEqual(results['scale'], (50, 32)) self.assertTupleEqual(results['scale_factor'], (50 / 25, 32 / 16)) # test min_width transform = RescaleToHeight(height=32, min_width=60) results = transform(copy.deepcopy(data_info)) self.assertTupleEqual(results['img'].shape[:2], (32, 60)) self.assertTupleEqual(results['scale'], (60, 32)) self.assertTupleEqual(results['scale_factor'], (60 / 25, 32 / 16)) # test max_width transform = RescaleToHeight(height=32, max_width=45) results = transform(copy.deepcopy(data_info)) self.assertTupleEqual(results['img'].shape[:2], (32, 45)) self.assertTupleEqual(results['scale'], (45, 32)) self.assertTupleEqual(results['scale_factor'], (45 / 25, 32 / 16)) # test width_divisor transform = RescaleToHeight(height=32, width_divisor=4) results = transform(copy.deepcopy(data_info)) self.assertTupleEqual(results['img'].shape[:2], (32, 48)) self.assertTupleEqual(results['scale'], (48, 32)) self.assertTupleEqual(results['scale_factor'], (48 / 25, 32 / 16)) def test_repr(self): transform = RescaleToHeight(height=32) self.assertEqual( repr(transform), ('RescaleToHeight(height=32, ' 'min_width=None, max_width=None, ' 'width_divisor=1, ' "resize_cfg={'type': 'Resize'})")) class TestPadToWidth(unittest.TestCase): def test_pad_to_width(self): data_info = dict(img=np.random.random((16, 25, 3))) # test size and size_divisor are both set with self.assertRaises(AssertionError): PadToWidth(width=10.5) transform = PadToWidth(width=100) results = transform(copy.deepcopy(data_info)) self.assertTupleEqual(results['img'].shape[:2], (16, 100)) self.assertEqual(results['valid_ratio'], 25 / 100) def test_repr(self): transform = PadToWidth(width=100) self.assertEqual( repr(transform), ("PadToWidth(width=100, pad_cfg={'type': 'Pad'})")) class TestSourceImagePad(unittest.TestCase): def setUp(self): img = np.zeros((15, 30, 3)) polygon = [np.array([10., 5., 20., 5., 20., 10., 10., 10.])] bbox = np.array([[10., 5., 20., 10.]]) self.data_info = dict(img=img, gt_polygons=polygon, gt_bboxes=bbox) def test_source_image_pad(self): # test image size equals to target size trans = SourceImagePad(target_scale=(30, 15)) target_polygon = self.data_info['gt_polygons'][0] target_bbox = self.data_info['gt_bboxes'] results = trans(self.data_info.copy()) self.assertEqual(results['img'].shape, (15, 30, 3)) self.assertEqual(results['img_shape'], (15, 30)) self.assertEqual(results['pad_shape'], (15, 30, 3)) self.assertEqual(results['pad_fixed_size'], (30, 15)) self.assertTrue(np.allclose(results['gt_polygons'][0], target_polygon)) self.assertTrue(np.allclose(results['gt_bboxes'][0], target_bbox)) # test pad to square trans = SourceImagePad(target_scale=30) target_polygon = np.array([10., 5., 20., 5., 20., 10., 10., 10.]) target_bbox = np.array([[10., 5., 20., 10.]]) results = trans(self.data_info.copy()) self.assertEqual(results['img'].shape, (30, 30, 3)) self.assertEqual(results['img_shape'], (30, 30)) self.assertEqual(results['pad_shape'], (30, 30, 3)) self.assertEqual(results['pad_fixed_size'], (30, 30)) self.assertTrue(np.allclose(results['gt_polygons'][0], target_polygon)) self.assertTrue(np.allclose(results['gt_bboxes'][0], target_bbox)) # test pad to different shape trans = SourceImagePad(target_scale=(40, 60)) target_polygon = np.array([10., 5., 20., 5., 20., 10., 10., 10.]) target_bbox = np.array([[10., 5., 20., 10.]]) results = trans(self.data_info.copy()) self.assertEqual(results['img'].shape, (60, 40, 3)) self.assertEqual(results['img_shape'], (60, 40)) self.assertEqual(results['pad_shape'], (60, 40, 3)) self.assertEqual(results['pad_fixed_size'], (40, 60)) self.assertTrue(np.allclose(results['gt_polygons'][0], target_polygon)) self.assertTrue(np.allclose(results['gt_bboxes'][0], target_bbox)) # test pad with different crop_ratio trans = SourceImagePad(target_scale=30, crop_ratio=1.0) target_polygon = np.array([10., 5., 20., 5., 20., 10., 10., 10.]) target_bbox = np.array([[10., 5., 20., 10.]]) results = trans(self.data_info.copy()) self.assertEqual(results['img'].shape, (30, 30, 3)) self.assertEqual(results['img_shape'], (30, 30)) self.assertEqual(results['pad_shape'], (30, 30, 3)) self.assertEqual(results['pad_fixed_size'], (30, 30)) self.assertTrue(np.allclose(results['gt_polygons'][0], target_polygon)) self.assertTrue(np.allclose(results['gt_bboxes'][0], target_bbox)) def test_repr(self): transform = SourceImagePad(target_scale=30, crop_ratio=0.1) self.assertEqual( repr(transform), ('SourceImagePad(target_scale = (30, 30), crop_ratio = (0.1, 0.1))' )) class TestRandomFlip(unittest.TestCase): def setUp(self): img = np.random.random((30, 40, 3)) gt_polygons = [np.array([10., 5., 20., 5., 20., 10., 10., 10.])] self.data_info = dict( img_shape=(30, 40), img=img, gt_polygons=gt_polygons) def test_flip_polygons(self): t = RandomFlip(prob=1.0, direction='horizontal') results = t.flip_polygons(self.data_info['gt_polygons'], (30, 40), 'horizontal') self.assertIsInstance(results, list) self.assertIsInstance(results[0], np.ndarray) self.assertTrue( (results[0] == np.array([30., 5., 20., 5., 20., 10., 30., 10.])).all()) results = t.flip_polygons(self.data_info['gt_polygons'], (30, 40), 'vertical') self.assertIsInstance(results, list) self.assertIsInstance(results[0], np.ndarray) self.assertTrue( (results[0] == np.array([10., 25., 20., 25., 20., 20., 10., 20.])).all()) results = t.flip_polygons(self.data_info['gt_polygons'], (30, 40), 'diagonal') self.assertIsInstance(results, list) self.assertIsInstance(results[0], np.ndarray) self.assertTrue( (results[0] == np.array([30., 25., 20., 25., 20., 20., 30., 20.])).all()) with self.assertRaises(ValueError): t.flip_polygons(self.data_info['gt_polygons'], (30, 40), 'mmocr') def test_flip(self): t = RandomFlip(prob=1.0, direction='horizontal') results = t(self.data_info.copy()) self.assertEqual(results['img'].shape, (30, 40, 3)) self.assertEqual(results['img_shape'], (30, 40)) self.assertTrue((results['gt_polygons'][0] == np.array( [30., 5., 20., 5., 20., 10., 30., 10.])).all()) class TestShortScaleAspectJitter(unittest.TestCase): @mock.patch('mmocr.datasets.pipelines.processing.np.random.random_sample') def test_transform(self, mock_random): ratio_range = (0.5, 1.5) aspect_ratio_range = (0.9, 1.1) mock_random.side_effect = [0.5, 0.5] img = np.zeros((15, 20, 3)) polygon = [np.array([10., 5., 20., 5., 20., 10., 10., 10.])] bbox = np.array([[10., 5., 20., 10.]]) data_info = dict(img=img, gt_polygons=polygon, gt_bboxes=bbox) t = ShortScaleAspectJitter( short_size=40, ratio_range=ratio_range, aspect_ratio_range=aspect_ratio_range, scale_divisor=4) results = t(data_info) self.assertEqual(results['img'].shape, (40, 56, 3)) self.assertEqual(results['img_shape'], (40, 56)) def test_repr(self): transform = ShortScaleAspectJitter( short_size=40, ratio_range=(0.5, 1.5), aspect_ratio_range=(0.9, 1.1), scale_divisor=4, resize_cfg=dict(type='Resize')) self.assertEqual( repr(transform), ('ShortScaleAspectJitter(' 'short_size = 40, ' 'ratio_range = (0.5, 1.5), ' 'aspect_ratio_range = (0.9, 1.1), ' 'scale_divisor = 4, ' "resize_cfg = {'type': 'Resize', 'scale': 0})")) class TestBoundedScaleAspectJitter(unittest.TestCase): @mock.patch('mmocr.datasets.pipelines.processing.np.random.random_sample') def test_transform(self, mock_random): mock_random.side_effect = [1.0, 1.0] data_info = dict(img=np.random.random((16, 25, 3)), img_shape=(16, 25)) # test size and size_divisor are both set transform = BoundedScaleAspectJitter(10, 5) result = transform(data_info) print(result['img'].shape) self.assertEqual(result['img'].shape, (8, 12, 3)) self.assertEqual(result['img_shape'], (8, 12)) def test_repr(self): transform = BoundedScaleAspectJitter(10, 5) print(repr(transform)) self.assertEqual( repr(transform), ('BoundedScaleAspectJitter(long_size_bound = 10, ' 'short_size_bound = 5, ratio_range = (0.7, 1.3), ' 'aspect_ratio_range = (0.9, 1.1), ' "resize_cfg = {'type': 'Resize', 'scale': 0})"))