Files
2026-07-13 12:46:08 +08:00

98 lines
3.1 KiB
Python

from __future__ import print_function, division, absolute_import
import sys
# unittest only added in 3.4 self.subTest()
if sys.version_info[0] < 3 or sys.version_info[1] < 4:
import unittest2 as unittest
else:
import unittest
# unittest.mock is not available in 2.7 (though unittest2 might contain it?)
try:
import unittest.mock as mock
except ImportError:
import mock
import numpy as np
from imgaug import augmenters as iaa
from imgaug.testutils import reseed, runtest_pickleable_uint8_img
class TestRandAugment(unittest.TestCase):
def setUp(self):
reseed()
# for some reason these mocks don't work with
# imgaug.augmenters.collections.(...)
@mock.patch("imgaug.augmenters.RandAugment._create_initial_augmenters_list")
@mock.patch("imgaug.augmenters.RandAugment._create_main_augmenters_list")
def test_n(self, mock_main, mock_initial):
mock_main.return_value = [iaa.Add(1), iaa.Add(2), iaa.Add(4)]
mock_initial.return_value = []
img = np.zeros((1, 1, 3), dtype=np.uint8)
expected = {
0: [0],
1: [1, 2, 4],
2: [1+1, 1+2, 1+4, 2+2, 2+4, 4+4]
}
for n in [0, 1, 2]:
with self.subTest(n=n):
aug = iaa.RandAugment(n=n)
img_aug = aug(image=img)
assert img_aug[0, 0, 0] in expected[n]
# for some reason these mocks don't work with
# imgaug.augmenters.collections.(...)
@mock.patch("imgaug.augmenters.RandAugment._create_initial_augmenters_list")
@mock.patch("imgaug.augmenters.RandAugment._create_main_augmenters_list")
def test_m(self, mock_main, mock_initial):
def _create_main_list(m, _cval):
return [iaa.Add(m)]
mock_main.side_effect = _create_main_list
mock_initial.return_value = []
img = np.zeros((1, 1, 3), dtype=np.uint8)
for m in [0, 1, 2]:
with self.subTest(m=m):
aug = iaa.RandAugment(m=m)
img_aug = aug(image=img)
assert img_aug[0, 0, 0] == m
def test_cval(self):
cval = 200
aug = iaa.RandAugment(n=1, m=30, cval=cval)
img = np.zeros((20, 20, 3), dtype=np.uint8)
x_cval = False
y_cval = False
# lots of iterations here, because only in some iterations an affine
# translation is actually applied
for _ in np.arange(500):
img_aug = aug(image=img)
x_cval = x_cval or np.all(img_aug[:, :1] == cval)
x_cval = x_cval or np.all(img_aug[:, -1:] == cval)
y_cval = y_cval or np.all(img_aug[:1, :] == cval)
y_cval = y_cval or np.all(img_aug[-1:, :] == cval)
if np.all([x_cval, y_cval]):
break
assert np.all([x_cval, y_cval])
def test_get_parameters(self):
aug = iaa.RandAugment(n=1, m=30, cval=100)
params = aug.get_parameters()
assert params[0] is aug[1].n
assert params[1] is aug._m
assert params[2] is aug._cval
def test_pickleable(self):
aug = iaa.RandAugment(m=(0, 10), n=(1, 2))
runtest_pickleable_uint8_img(aug, iterations=50)