242 lines
8.3 KiB
Python
242 lines
8.3 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
|
|
import cv2
|
|
|
|
import imgaug as ia
|
|
from imgaug import augmenters as iaa
|
|
from imgaug import random as iarandom
|
|
from imgaug.testutils import reseed, runtest_pickleable_uint8_img
|
|
import imgaug.augmenters.color as colorlib
|
|
import imgaug.augmenters.artistic as artisticlib
|
|
|
|
|
|
class Test_stylize_cartoon(unittest.TestCase):
|
|
@classmethod
|
|
def _test_integrationtest(cls, size, validate_grads):
|
|
image = ia.data.quokka_square((size, size))
|
|
|
|
image_cartoon = iaa.stylize_cartoon(image, blur_ksize=5,
|
|
segmentation_size=2.0)
|
|
|
|
image_avg = np.average(image.astype(np.float32), axis=2)
|
|
image_cartoon_avg = np.average(image_cartoon.astype(np.float32), axis=2)
|
|
|
|
if validate_grads:
|
|
gradx_image = image_avg[:, :-1] - image_avg[:, 1:]
|
|
grady_image = image_avg[:-1, :] - image_avg[1:, :]
|
|
gradx_cartoon = image_cartoon_avg[:, :-1] - image_cartoon_avg[:, 1:]
|
|
grady_cartoon = image_cartoon_avg[:-1, :] - image_cartoon_avg[1:, :]
|
|
|
|
assert (
|
|
(
|
|
np.average(np.abs(gradx_cartoon))
|
|
+ np.average(np.abs(grady_cartoon))
|
|
)
|
|
<
|
|
(
|
|
np.average(np.abs(gradx_image))
|
|
+ np.average(np.abs(grady_image))
|
|
)
|
|
)
|
|
|
|
# average saturation of cartoon image should be increased
|
|
image_hsv = colorlib.change_colorspace_(np.copy(image),
|
|
to_colorspace=iaa.CSPACE_HSV)
|
|
cartoon_hsv = colorlib.change_colorspace_(np.copy(image_cartoon),
|
|
to_colorspace=iaa.CSPACE_HSV)
|
|
assert (
|
|
np.average(cartoon_hsv[:, :, 1])
|
|
> np.average(image_hsv[:, :, 1])
|
|
)
|
|
|
|
# as edges are all drawn in completely black, there should be more
|
|
# completely black pixels in the cartoon image
|
|
image_black = np.sum(image_avg <= 0.01)
|
|
cartoon_black = np.sum(image_cartoon_avg <= 0.01)
|
|
|
|
assert cartoon_black > image_black
|
|
|
|
def test_integrationtest(self):
|
|
self._test_integrationtest(128, True)
|
|
|
|
def test_integrationtest_large_image(self):
|
|
# TODO the validation of gradients currently doesn't work well
|
|
# for the laplacian edge method, but it should
|
|
self._test_integrationtest(400, False)
|
|
|
|
@mock.patch("cv2.medianBlur")
|
|
def test_blur_ksize_is_1(self, mock_blur):
|
|
def _side_effect(image, ksize):
|
|
return image
|
|
|
|
mock_blur.side_effect = _side_effect
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
|
|
_ = iaa.stylize_cartoon(image, blur_ksize=1)
|
|
|
|
# median blur is called another time in _find_edge_laplacian, but
|
|
# that function is only called if the image is larger
|
|
assert mock_blur.call_count == 0
|
|
|
|
@mock.patch("cv2.medianBlur")
|
|
def test_blur_ksize_gt_1(self, mock_blur):
|
|
def _side_effect(image, ksize):
|
|
return image
|
|
|
|
mock_blur.side_effect = _side_effect
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
|
|
_ = iaa.stylize_cartoon(image, blur_ksize=7)
|
|
|
|
assert mock_blur.call_count == 1
|
|
assert mock_blur.call_args_list[0][0][1] == 7
|
|
|
|
@mock.patch("cv2.pyrMeanShiftFiltering")
|
|
def test_segmentation_size_is_0(self, mock_msf):
|
|
def _side_effect(image, sp, sr, dst):
|
|
dst[...] = image
|
|
|
|
mock_msf.side_effect = _side_effect
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
|
|
_ = iaa.stylize_cartoon(image, segmentation_size=0.0)
|
|
|
|
assert mock_msf.call_count == 0
|
|
|
|
@mock.patch("cv2.pyrMeanShiftFiltering")
|
|
def test_segmentation_size_gt_0(self, mock_msf):
|
|
def _side_effect(image, sp, sr, dst):
|
|
dst[...] = image
|
|
|
|
mock_msf.side_effect = _side_effect
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
|
|
_ = iaa.stylize_cartoon(image, segmentation_size=0.5)
|
|
|
|
assert mock_msf.call_count == 1
|
|
assert np.allclose(mock_msf.call_args_list[0][1]["sp"], 10*0.5)
|
|
assert np.allclose(mock_msf.call_args_list[0][1]["sr"], 20*0.5)
|
|
|
|
@mock.patch("imgaug.augmenters.artistic._suppress_edge_blobs")
|
|
def test_suppress_edges_true(self, mock_seb):
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
mock_seb.return_value = np.copy(image[..., 0])
|
|
|
|
_ = iaa.stylize_cartoon(image, suppress_edges=True)
|
|
|
|
assert mock_seb.call_count == 2
|
|
|
|
@mock.patch("imgaug.augmenters.artistic._suppress_edge_blobs")
|
|
def test_suppress_edges_false(self, mock_seb):
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
|
|
_ = iaa.stylize_cartoon(image, suppress_edges=False)
|
|
|
|
assert mock_seb.call_count == 0
|
|
|
|
@mock.patch("imgaug.augmenters.artistic._find_edges_laplacian")
|
|
def test_large_image(self, mock_fel):
|
|
def _side_effect_fel(image, edge_multiplier, from_colorspace):
|
|
return image[..., 0]
|
|
|
|
mock_fel.side_effect = _side_effect_fel
|
|
image = np.zeros((10, 401, 3), dtype=np.uint8)
|
|
|
|
_ = iaa.stylize_cartoon(image, segmentation_size=0)
|
|
|
|
assert mock_fel.call_count == 1
|
|
|
|
|
|
class Test__saturate(unittest.TestCase):
|
|
def _get_avg_saturation(self, image):
|
|
hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
|
|
return np.average(hsv[..., 1])
|
|
|
|
def test_saturation_is_1(self):
|
|
image = np.array([
|
|
[10, 20, 30],
|
|
[40, 50, 60],
|
|
[70, 80, 90],
|
|
[100, 110, 120],
|
|
[10, 10, 10],
|
|
[100, 0, 0],
|
|
[0, 100, 0],
|
|
[0, 0, 100]
|
|
], dtype=np.uint8).reshape((1, 8, 3))
|
|
|
|
observed_1 = artisticlib._saturate(image, 1.0, colorlib.CSPACE_RGB)
|
|
observed_2 = artisticlib._saturate(image, 2.0, colorlib.CSPACE_RGB)
|
|
|
|
sat_img = self._get_avg_saturation(image)
|
|
sat_1 = self._get_avg_saturation(observed_1)
|
|
sat_2 = self._get_avg_saturation(observed_2)
|
|
assert sat_img < sat_2
|
|
assert sat_1 < sat_2
|
|
|
|
|
|
class TestCartoon(unittest.TestCase):
|
|
def setUp(self):
|
|
reseed()
|
|
|
|
def test___init___defaults(self):
|
|
aug = iaa.Cartoon()
|
|
assert aug.blur_ksize.a.value == 1
|
|
assert aug.blur_ksize.b.value == 5
|
|
assert np.isclose(aug.segmentation_size.a.value, 0.8)
|
|
assert np.isclose(aug.segmentation_size.b.value, 1.2)
|
|
assert np.isclose(aug.edge_prevalence.a.value, 0.9)
|
|
assert np.isclose(aug.edge_prevalence.b.value, 1.1)
|
|
assert np.isclose(aug.saturation.a.value, 1.5)
|
|
assert np.isclose(aug.saturation.b.value, 2.5)
|
|
assert aug.from_colorspace == iaa.CSPACE_RGB
|
|
|
|
def test_draw_samples(self):
|
|
mock_batch = mock.Mock()
|
|
mock_batch.nb_rows = 50
|
|
aug = iaa.Cartoon()
|
|
rs = iarandom.RNG(0)
|
|
|
|
samples = aug._draw_samples(mock_batch, rs)
|
|
|
|
assert len(np.unique(np.round(samples[0]*100, decimals=0))) > 1
|
|
assert len(np.unique(np.round(samples[1]*100, decimals=0))) > 1
|
|
assert len(np.unique(np.round(samples[2]*100, decimals=0))) > 1
|
|
assert len(np.unique(np.round(samples[3]*100, decimals=0))) > 1
|
|
|
|
@mock.patch("imgaug.augmenters.artistic.stylize_cartoon")
|
|
def test_call_of_stylize_cartoon(self, mock_sc):
|
|
image = np.arange(4*4*3).astype(np.uint8).reshape((4, 4, 3))
|
|
aug = iaa.Cartoon()
|
|
|
|
mock_sc.return_value = np.copy(image)
|
|
|
|
_ = aug(images=[image, image])
|
|
|
|
assert mock_sc.call_count == 2
|
|
|
|
def test_get_parameters(self):
|
|
aug = iaa.Cartoon()
|
|
params = aug.get_parameters()
|
|
assert params[0] is aug.blur_ksize
|
|
assert params[1] is aug.segmentation_size
|
|
assert params[2] is aug.saturation
|
|
assert params[3] is aug.edge_prevalence
|
|
assert params[4] == iaa.CSPACE_RGB
|
|
|
|
def test_pickleable(self):
|
|
aug = iaa.Cartoon(seed=1)
|
|
runtest_pickleable_uint8_img(aug, iterations=6)
|