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

515 lines
16 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 six.moves as sm
import imgaug as ia
# TODO add tests for:
# hooks is_activated
# hooks is_propagating
# hooks preprocess
# hooks postprocess
# HeatmapsOnImage.__init__()
# HeatmapsOnImage.get_arr()
# HeatmapsOnImage.to_uint8()
# HeatmapsOnImage.from_0to1()
# HeatmapsOnImage.copy()
# HeatmapsOnImage.deepcopy()
class TestHeatmapsOnImage_draw(unittest.TestCase):
def test_basic_functionality(self):
heatmaps_arr = np.float32([
[0.5, 0.0, 0.0, 0.5],
[0.0, 1.0, 1.0, 0.0],
[0.0, 1.0, 1.0, 0.0],
[0.5, 0.0, 0.0, 0.5],
])
heatmaps = ia.HeatmapsOnImage(heatmaps_arr, shape=(4, 4, 3))
heatmaps_drawn = heatmaps.draw()[0]
assert heatmaps_drawn.shape == (4, 4, 3)
v1 = heatmaps_drawn[0, 1]
v2 = heatmaps_drawn[0, 0]
v3 = heatmaps_drawn[1, 1]
v1_coords = [(0, 1), (0, 2), (1, 0), (1, 3), (2, 0), (2, 3), (3, 1),
(3, 2)]
v2_coords = [(0, 0), (0, 3), (3, 0), (3, 3)]
v3_coords = [(1, 1), (1, 2), (2, 1), (2, 2)]
for y, x in v1_coords:
assert np.allclose(heatmaps_drawn[y, x], v1)
for y, x in v2_coords:
assert np.allclose(heatmaps_drawn[y, x], v2)
for y, x in v3_coords:
assert np.allclose(heatmaps_drawn[y, x], v3)
def test_use_size_arg_with_different_shape_than_heatmap_arr_shape(self):
# size differs from heatmap array size
heatmaps_arr = np.float32([
[0.0, 1.0],
[0.0, 1.0]
])
heatmaps = ia.HeatmapsOnImage(heatmaps_arr, shape=(2, 2, 3))
heatmaps_drawn = heatmaps.draw(size=(4, 4))[0]
assert heatmaps_drawn.shape == (4, 4, 3)
v1 = heatmaps_drawn[0, 0]
v2 = heatmaps_drawn[0, -1]
for y in sm.xrange(4):
for x in sm.xrange(2):
assert np.allclose(heatmaps_drawn[y, x], v1)
for y in sm.xrange(4):
for x in sm.xrange(2, 4):
assert np.allclose(heatmaps_drawn[y, x], v2)
# TODO test other cmaps
class TestHeatmapsOnImage_draw_on_image(unittest.TestCase):
@property
def heatmaps(self):
heatmaps_arr = np.float32([
[0.0, 1.0],
[0.0, 1.0]
])
return ia.HeatmapsOnImage(heatmaps_arr, shape=(2, 2, 3))
def test_cmap_is_none(self):
heatmaps = self.heatmaps
image = np.uint8([
[0, 0, 0, 255],
[0, 0, 0, 255],
[0, 0, 0, 255],
[0, 0, 0, 255]
])
image = np.tile(image[..., np.newaxis], (1, 1, 3))
heatmaps_drawn = heatmaps.draw_on_image(image, alpha=0.5, cmap=None)[0]
assert heatmaps_drawn.shape == (4, 4, 3)
assert np.all(heatmaps_drawn[0:4, 0:2, :] == 0)
assert (
np.all(heatmaps_drawn[0:4, 2:3, :] == 128)
or np.all(heatmaps_drawn[0:4, 2:3, :] == 127))
assert (
np.all(heatmaps_drawn[0:4, 3:4, :] == 255)
or np.all(heatmaps_drawn[0:4, 3:4, :] == 254))
def test_cmap_is_none_and_resize_is_image(self):
heatmaps = self.heatmaps
image = np.uint8([
[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]
])
image = np.tile(image[..., np.newaxis], (1, 1, 3))
heatmaps_drawn = heatmaps.draw_on_image(
image, alpha=0.5, resize="image", cmap=None)[0]
assert heatmaps_drawn.shape == (2, 2, 3)
assert np.all(heatmaps_drawn[0:2, 0, :] == 0)
assert (
np.all(heatmaps_drawn[0:2, 1, :] == 128)
or np.all(heatmaps_drawn[0:2, 1, :] == 127))
class TestHeatmapsOnImage_invert(unittest.TestCase):
@property
def heatmaps_arr(self):
return np.float32([
[0.0, 5.0, 10.0],
[-1.0, -2.0, 7.5]
])
@property
def expected_arr(self):
return np.float32([
[8.0, 3.0, -2.0],
[9.0, 10.0, 0.5]
])
def test_with_2d_input_array(self):
# (H, W)
heatmaps_arr = self.heatmaps_arr
expected = self.expected_arr
heatmaps = ia.HeatmapsOnImage(heatmaps_arr,
shape=(2, 3),
min_value=-2.0,
max_value=10.0)
assert np.allclose(heatmaps.get_arr(), heatmaps_arr)
assert np.allclose(heatmaps.invert().get_arr(), expected)
def test_with_3d_input_array(self):
# (H, W, 1)
heatmaps_arr = self.heatmaps_arr
expected = self.expected_arr
heatmaps = ia.HeatmapsOnImage(heatmaps_arr[..., np.newaxis],
shape=(2, 3),
min_value=-2.0,
max_value=10.0)
assert np.allclose(heatmaps.get_arr(),
heatmaps_arr[..., np.newaxis])
assert np.allclose(heatmaps.invert().get_arr(),
expected[..., np.newaxis])
class TestHeatmapsOnImage_pad(unittest.TestCase):
@property
def heatmaps(self):
heatmaps_arr = np.float32([
[0.0, 1.0],
[0.0, 1.0]
])
return ia.HeatmapsOnImage(heatmaps_arr, shape=(2, 2, 3))
def test_defaults(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad(top=1, right=2, bottom=3, left=4)
assert heatmaps_padded.arr_0to1.shape == (2+(1+3), 2+(4+2), 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
])
)
def test_mode_constant_with_cval_050(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad(top=1, right=2, bottom=3, left=4,
cval=0.5)
assert heatmaps_padded.arr_0to1.shape == (2+(1+3), 2+(4+2), 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
[0.5, 0.5, 0.5, 0.5, 0.0, 1.0, 0.5, 0.5],
[0.5, 0.5, 0.5, 0.5, 0.0, 1.0, 0.5, 0.5],
[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
])
)
def test_mode_edge(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad(top=1, right=2, bottom=3, left=4,
mode="edge")
assert heatmaps_padded.arr_0to1.shape == (2+(1+3), 2+(4+2), 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]
])
)
class TestHeatmapsOnImage_pad_to_aspect_ratio(unittest.TestCase):
@property
def heatmaps(self):
heatmaps_arr = np.float32([
[0.0, 0.0, 1.0],
[0.0, 0.0, 1.0]
])
return ia.HeatmapsOnImage(heatmaps_arr, shape=(2, 2, 3))
def test_square_ratio_with_default_mode_and_cval(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad_to_aspect_ratio(1.0)
assert heatmaps_padded.arr_0to1.shape == (3, 3, 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 1.0],
[0.0, 0.0, 1.0],
[0.0, 0.0, 0.0]
])
)
def test_square_ratio_with_cval_050(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad_to_aspect_ratio(1.0, cval=0.5)
assert heatmaps_padded.arr_0to1.shape == (3, 3, 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 1.0],
[0.0, 0.0, 1.0],
[0.5, 0.5, 0.5]
])
)
def test_square_ratio_with_edge_mode(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad_to_aspect_ratio(1.0, mode="edge")
assert heatmaps_padded.arr_0to1.shape == (3, 3, 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 1.0],
[0.0, 0.0, 1.0],
[0.0, 0.0, 1.0]
])
)
def test_wider_than_high_ratio_with_cval_010(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad_to_aspect_ratio(2.0, cval=0.1)
assert heatmaps_padded.arr_0to1.shape == (2, 4, 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 1.0, 0.1],
[0.0, 0.0, 1.0, 0.1]
])
)
def test_higher_than_wide_ratio_with_cval_010(self):
heatmaps = self.heatmaps
heatmaps_padded = heatmaps.pad_to_aspect_ratio(0.25, cval=0.1)
assert heatmaps_padded.arr_0to1.shape == (12, 3, 1)
assert np.allclose(
heatmaps_padded.arr_0to1[:, :, 0],
np.float32([
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.0, 0.0, 1.0],
[0.0, 0.0, 1.0],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1],
[0.1, 0.1, 0.1]
])
)
class TestHeatmapsOnImage_avg_pool(unittest.TestCase):
def test_with_kernel_size_2(self):
heatmaps_arr = np.float32([
[0.0, 0.0, 0.5, 1.0],
[0.0, 0.0, 0.5, 1.0],
[0.0, 0.0, 0.5, 1.0],
[0.0, 0.0, 0.5, 1.0]
])
heatmaps = ia.HeatmapsOnImage(heatmaps_arr, shape=(4, 4, 3))
heatmaps_pooled = heatmaps.avg_pool(2)
assert heatmaps_pooled.arr_0to1.shape == (2, 2, 1)
assert np.allclose(
heatmaps_pooled.arr_0to1[:, :, 0],
np.float32([[0.0, 0.75],
[0.0, 0.75]])
)
class TestHeatmapsOnImage_max_pool(unittest.TestCase):
def test_with_kernel_size_2(self):
heatmaps_arr = np.float32([
[0.0, 0.0, 0.5, 1.0],
[0.0, 0.0, 0.5, 1.0],
[0.0, 0.0, 0.5, 1.0],
[0.0, 0.0, 0.5, 1.0]
])
heatmaps = ia.HeatmapsOnImage(heatmaps_arr, shape=(4, 4, 3))
heatmaps_pooled = heatmaps.max_pool(2)
assert heatmaps_pooled.arr_0to1.shape == (2, 2, 1)
assert np.allclose(
heatmaps_pooled.arr_0to1[:, :, 0],
np.float32([[0.0, 1.0],
[0.0, 1.0]])
)
class TestHeatmapsOnImage_resize(unittest.TestCase):
def test_resize_to_exact_shape(self):
heatmaps_arr = np.float32([
[0.0, 1.0]
])
heatmaps = ia.HeatmapsOnImage(heatmaps_arr, shape=(4, 4, 3))
heatmaps_scaled = heatmaps.resize((4, 4), interpolation="nearest")
assert heatmaps_scaled.arr_0to1.shape == (4, 4, 1)
assert heatmaps_scaled.arr_0to1.dtype.name == "float32"
assert np.allclose(
heatmaps_scaled.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 1.0, 1.0],
[0.0, 0.0, 1.0, 1.0],
[0.0, 0.0, 1.0, 1.0],
[0.0, 0.0, 1.0, 1.0]
])
)
def test_resize_to_twice_the_size(self):
heatmaps_arr = np.float32([
[0.0, 1.0]
])
heatmaps = ia.HeatmapsOnImage(heatmaps_arr, shape=(4, 4, 3))
heatmaps_scaled = heatmaps.resize(2.0, interpolation="nearest")
assert heatmaps_scaled.arr_0to1.shape == (2, 4, 1)
assert heatmaps_scaled.arr_0to1.dtype.name == "float32"
assert np.allclose(
heatmaps_scaled.arr_0to1[:, :, 0],
np.float32([
[0.0, 0.0, 1.0, 1.0],
[0.0, 0.0, 1.0, 1.0]
])
)
class TestHeatmapsOnImage_from_uint8(unittest.TestCase):
def test_3d_uint8_array(self):
hm = ia.HeatmapsOnImage.from_uint8(
np.uint8([
[0, 128, 255],
[255, 128, 0]
])[..., np.newaxis],
(20, 30, 3)
)
assert hm.shape == (20, 30, 3)
assert hm.arr_0to1.shape == (2, 3, 1)
assert np.allclose(hm.arr_0to1[..., 0], np.float32([
[0, 128/255, 1.0],
[1.0, 128/255, 0]
]))
def test_2d_uint8_array(self):
hm = ia.HeatmapsOnImage.from_uint8(
np.uint8([
[0, 128, 255],
[255, 128, 0]
]),
(20, 30, 3)
)
assert hm.shape == (20, 30, 3)
assert hm.arr_0to1.shape == (2, 3, 1)
assert np.allclose(hm.arr_0to1[..., 0], np.float32([
[0, 128/255, 1.0],
[1.0, 128/255, 0]
]))
def test_min_value_and_max_value(self):
# min_value, max_value
hm = ia.HeatmapsOnImage.from_uint8(
np.uint8([
[0, 128, 255],
[255, 128, 0]
])[..., np.newaxis],
(20, 30, 3),
min_value=-1.0,
max_value=2.0
)
assert hm.shape == (20, 30, 3)
assert hm.arr_0to1.shape == (2, 3, 1)
assert np.allclose(hm.arr_0to1[..., 0], np.float32([
[0, 128/255, 1.0],
[1.0, 128/255, 0]
]))
assert np.allclose(hm.min_value, -1.0)
assert np.allclose(hm.max_value, 2.0)
class TestHeatmapsOnImage_change_normalization(unittest.TestCase):
def test_increase_max_value(self):
# (0.0, 1.0) -> (0.0, 2.0)
arr = np.float32([
[0.0, 0.5, 1.0],
[1.0, 0.5, 0.0]
])
observed = ia.HeatmapsOnImage.change_normalization(
arr, (0.0, 1.0), (0.0, 2.0))
expected = np.float32([
[0.0, 1.0, 2.0],
[2.0, 1.0, 0.0]
])
assert np.allclose(observed, expected)
def test_decrease_min_and_max_value(self):
# (0.0, 1.0) -> (-1.0, 0.0)
arr = np.float32([
[0.0, 0.5, 1.0],
[1.0, 0.5, 0.0]
])
observed = ia.HeatmapsOnImage.change_normalization(
arr, (0.0, 1.0), (-1.0, 0.0))
expected = np.float32([
[-1.0, -0.5, 0.0],
[0.0, -0.5, -1.0]
])
assert np.allclose(observed, expected)
def test_increase_min_and_max_value__non_standard_source(self):
# (-1.0, 1.0) -> (1.0, 3.0)
arr = np.float32([
[-1.0, 0.0, 1.0],
[1.0, 0.0, -1.0]
])
observed = ia.HeatmapsOnImage.change_normalization(
arr, (-1.0, 1.0), (1.0, 3.0))
expected = np.float32([
[1.0, 2.0, 3.0],
[3.0, 2.0, 1.0]
])
assert np.allclose(observed, expected)
def test_value_ranges_given_as_heatmaps_on_image(self):
# (-1.0, 1.0) -> (1.0, 3.0)
# value ranges given as HeatmapsOnImage
arr = np.float32([
[-1.0, 0.0, 1.0],
[1.0, 0.0, -1.0]
])
source = ia.HeatmapsOnImage(
np.float32([[0.0]]), min_value=-1.0, max_value=1.0, shape=(1, 1, 3))
target = ia.HeatmapsOnImage(
np.float32([[1.0]]), min_value=1.0, max_value=3.0, shape=(1, 1, 3))
observed = ia.HeatmapsOnImage.change_normalization(arr, source, target)
expected = np.float32([
[1.0, 2.0, 3.0],
[3.0, 2.0, 1.0]
])
assert np.allclose(observed, expected)