forked from PaddlePaddle/PaddleOCR
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_iaa_augment.py
181 lines (147 loc) · 6.65 KB
/
test_iaa_augment.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
import os
import sys
import pytest
import numpy as np
import random
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.abspath(os.path.join(current_dir, "..")))
from ppocr.data.imaug.iaa_augment import IaaAugment
# Set a fixed random seed to ensure test reproducibility
np.random.seed(42)
random.seed(42)
# Fixture to provide a sample image for tests
@pytest.fixture
def sample_image():
# Create a 100x100 pixel dummy image with 3 color channels (RGB)
return np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
# Fixture to provide sample polygons for tests
@pytest.fixture
def sample_polys():
# Create dummy polygons as sample data
polys = [
np.array([[10, 10], [20, 10], [20, 20], [10, 20]], dtype=np.float32),
np.array([[30, 30], [40, 30], [40, 40], [30, 40]], dtype=np.float32),
]
return polys
# Helper function to create a data dictionary for testing
def create_data(sample_image, sample_polys):
return {
"image": sample_image.copy(),
"polys": [poly.copy() for poly in sample_polys],
}
# Test the default behavior of the augmenter (without specified arguments)
def test_iaa_augment_default(sample_image, sample_polys):
data = create_data(sample_image, sample_polys)
augmenter = IaaAugment()
transformed_data = augmenter(data)
# Check the data types and structure of the transformed image and polygons
assert isinstance(
transformed_data["image"], np.ndarray
), "Image should be a numpy array"
assert isinstance(
transformed_data["polys"], np.ndarray
), "Polys should be a numpy array"
assert transformed_data["image"].ndim == 3, "Image should be 3-dimensional"
# Verify that the polygons have been transformed
polys_changed = any(
not np.allclose(orig_poly, trans_poly)
for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
)
assert polys_changed, "Polygons should have been transformed"
# Test the augmenter with empty arguments, meaning no transformations should occur
def test_iaa_augment_none(sample_image, sample_polys):
data = create_data(sample_image, sample_polys)
augmenter = IaaAugment(augmenter_args=[])
transformed_data = augmenter(data)
# Check that the image and polygons remain unchanged
assert np.array_equal(
data["image"], transformed_data["image"]
), "Image should be unchanged"
for orig_poly, transformed_poly in zip(data["polys"], transformed_data["polys"]):
assert np.array_equal(
orig_poly, transformed_poly
), "Polygons should be unchanged"
# Parameterized test to check various augmenter arguments and expected image shapes
@pytest.mark.parametrize(
"augmenter_args, expected_shape",
[
([], (100, 100, 3)),
([{"type": "Resize", "args": {"size": [0.5, 0.5]}}], (50, 50, 3)),
([{"type": "Resize", "args": {"size": [2.0, 2.0]}}], (200, 200, 3)),
],
)
def test_iaa_augment_resize(sample_image, sample_polys, augmenter_args, expected_shape):
data = create_data(sample_image, sample_polys)
augmenter = IaaAugment(augmenter_args=augmenter_args)
transformed_data = augmenter(data)
# Verify that the transformed image has the expected shape
assert (
transformed_data["image"].shape == expected_shape
), f"Expected image shape {expected_shape}, got {transformed_data['image'].shape}"
# Test custom augmenter arguments with specific transformations
def test_iaa_augment_custom(sample_image, sample_polys):
data = create_data(sample_image, sample_polys)
augmenter_args = [
{"type": "Affine", "args": {"rotate": [45, 45]}}, # Apply 45-degree rotation
{"type": "Resize", "args": {"size": [0.5, 0.5]}},
]
augmenter = IaaAugment(augmenter_args=augmenter_args)
transformed_data = augmenter(data)
# Check the expected image dimensions after resizing
expected_height = int(sample_image.shape[0] * 0.5)
expected_width = int(sample_image.shape[1] * 0.5)
assert (
transformed_data["image"].shape[0] == expected_height
), "Image height should be scaled by 0.5"
assert (
transformed_data["image"].shape[1] == expected_width
), "Image width should be scaled by 0.5"
# Verify that the polygons have been transformed
polys_changed = any(
not np.allclose(orig_poly, trans_poly)
for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
)
assert polys_changed, "Polygons should have been transformed"
# Test that an unknown transformation type raises an AttributeError
def test_iaa_augment_unknown_transform():
augmenter_args = [{"type": "UnknownTransform", "args": {}}]
with pytest.raises(AttributeError):
IaaAugment(augmenter_args=augmenter_args)
# Test that an invalid resize size parameter raises a ValueError
def test_iaa_augment_invalid_resize_size(sample_image, sample_polys):
augmenter_args = [{"type": "Resize", "args": {"size": "invalid_size"}}]
with pytest.raises(ValueError) as exc_info:
IaaAugment(augmenter_args=augmenter_args)
assert "'size' must be a list or tuple of two numbers" in str(exc_info.value)
# Test that polygons are transformed as expected
def test_iaa_augment_polys_transformation(sample_image, sample_polys):
data = create_data(sample_image, sample_polys)
augmenter_args = [
{"type": "Affine", "args": {"rotate": [90, 90]}}, # Apply 90-degree rotation
]
augmenter = IaaAugment(augmenter_args=augmenter_args)
transformed_data = augmenter(data)
# Verify that the polygons have been transformed
polys_changed = any(
not np.allclose(orig_poly, trans_poly)
for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
)
assert polys_changed, "Polygons should have been transformed"
# Test multiple transformations applied to the augmenter
def test_iaa_augment_multiple_transforms(sample_image, sample_polys):
augmenter_args = [
{"type": "Fliplr", "args": {"p": 1.0}}, # Always apply horizontal flip
{"type": "Affine", "args": {"shear": 10}},
]
data = create_data(sample_image, sample_polys)
augmenter = IaaAugment(augmenter_args=augmenter_args)
transformed_data = augmenter(data)
# Ensure the image has been transformed
images_different = not np.array_equal(transformed_data["image"], sample_image)
assert images_different, "Image should be transformed"
# Ensure the polygons have been transformed
polys_changed = any(
not np.allclose(orig_poly, trans_poly)
for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
)
assert polys_changed, "Polygons should have been transformed"