| |
| import math |
| import unittest.mock as mock |
|
|
| import numpy as np |
| import torch |
| import torchvision.transforms.functional as TF |
| from PIL import Image |
|
|
| import mmocr.datasets.pipelines.ocr_transforms as transforms |
|
|
|
|
| def test_resize_ocr(): |
| input_img = np.ones((64, 256, 3), dtype=np.uint8) |
|
|
| rci = transforms.ResizeOCR( |
| 32, min_width=32, max_width=160, keep_aspect_ratio=True) |
| results = {'img_shape': input_img.shape, 'img': input_img} |
|
|
| |
| results = rci(results) |
| assert np.allclose([32, 160, 3], results['pad_shape']) |
| assert np.allclose([32, 160, 3], results['img'].shape) |
| assert 'valid_ratio' in results |
| assert math.isclose(results['valid_ratio'], 0.8) |
| assert math.isclose(np.sum(results['img'][:, 129:, :]), 0) |
|
|
| rci = transforms.ResizeOCR( |
| 32, min_width=32, max_width=160, keep_aspect_ratio=False) |
| results = {'img_shape': input_img.shape, 'img': input_img} |
| results = rci(results) |
| assert math.isclose(results['valid_ratio'], 1) |
|
|
|
|
| def test_to_tensor(): |
| input_img = np.ones((64, 256, 3), dtype=np.uint8) |
|
|
| expect_output = TF.to_tensor(input_img) |
| rci = transforms.ToTensorOCR() |
|
|
| results = {'img': input_img} |
| results = rci(results) |
|
|
| assert np.allclose(results['img'].numpy(), expect_output.numpy()) |
|
|
|
|
| def test_normalize(): |
| inputs = torch.zeros(3, 10, 10) |
|
|
| expect_output = torch.ones_like(inputs) * (-1) |
| rci = transforms.NormalizeOCR(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
|
|
| results = {'img': inputs} |
| results = rci(results) |
|
|
| assert np.allclose(results['img'].numpy(), expect_output.numpy()) |
|
|
|
|
| @mock.patch('%s.transforms.np.random.random' % __name__) |
| def test_online_crop(mock_random): |
| kwargs = dict( |
| box_keys=['x1', 'y1', 'x2', 'y2', 'x3', 'y3', 'x4', 'y4'], |
| jitter_prob=0.5, |
| max_jitter_ratio_x=0.05, |
| max_jitter_ratio_y=0.02) |
|
|
| mock_random.side_effect = [0.1, 1, 1, 1] |
|
|
| src_img = np.ones((100, 100, 3), dtype=np.uint8) |
| results = { |
| 'img': src_img, |
| 'img_info': { |
| 'x1': '20', |
| 'y1': '20', |
| 'x2': '40', |
| 'y2': '20', |
| 'x3': '40', |
| 'y3': '40', |
| 'x4': '20', |
| 'y4': '40' |
| } |
| } |
|
|
| rci = transforms.OnlineCropOCR(**kwargs) |
|
|
| results = rci(results) |
|
|
| assert np.allclose(results['img_shape'], [20, 20, 3]) |
|
|
| |
| mock_random.side_effect = [0.1, 1, 1, 1] |
| results['img_info'] = {} |
| results['img'] = src_img |
|
|
| results = rci(results) |
| assert np.allclose(results['img'].shape, [100, 100, 3]) |
|
|
|
|
| def test_fancy_pca(): |
| input_tensor = torch.rand(3, 32, 100) |
|
|
| rci = transforms.FancyPCA() |
|
|
| results = {'img': input_tensor} |
| results = rci(results) |
|
|
| assert results['img'].shape == torch.Size([3, 32, 100]) |
|
|
|
|
| @mock.patch('%s.transforms.np.random.uniform' % __name__) |
| def test_random_padding(mock_random): |
| kwargs = dict(max_ratio=[0.0, 0.0, 0.0, 0.0], box_type=None) |
|
|
| mock_random.side_effect = [1, 1, 1, 1] |
|
|
| src_img = np.ones((32, 100, 3), dtype=np.uint8) |
| results = {'img': src_img, 'img_shape': (32, 100, 3)} |
|
|
| rci = transforms.RandomPaddingOCR(**kwargs) |
|
|
| results = rci(results) |
| print(results['img'].shape) |
| assert np.allclose(results['img_shape'], [96, 300, 3]) |
|
|
|
|
| def test_opencv2pil(): |
| src_img = np.ones((32, 100, 3), dtype=np.uint8) |
| results = {'img': src_img} |
| rci = transforms.OpencvToPil() |
|
|
| results = rci(results) |
| assert np.allclose(results['img'].size, (100, 32)) |
|
|
|
|
| def test_pil2opencv(): |
| src_img = Image.new('RGB', (100, 32), color=(255, 255, 255)) |
| results = {'img': src_img} |
| rci = transforms.PilToOpencv() |
|
|
| results = rci(results) |
| assert np.allclose(results['img'].shape, (32, 100, 3)) |
|
|
