983 lines
39 KiB
Python
983 lines
39 KiB
Python
from __future__ import print_function, division, absolute_import
|
|
|
|
import time
|
|
import multiprocessing
|
|
import pickle
|
|
from collections import defaultdict
|
|
import warnings
|
|
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
|
|
import imgaug.multicore as multicore
|
|
import imgaug.random as iarandom
|
|
from imgaug import augmenters as iaa
|
|
from imgaug.testutils import reseed
|
|
from imgaug.augmentables.batches import Batch, UnnormalizedBatch
|
|
|
|
IS_SUPPORTING_CONTEXTS = (sys.version_info[0] == 3
|
|
and sys.version_info[1] >= 4)
|
|
|
|
|
|
class clean_context():
|
|
def __init__(self):
|
|
self.old_context = None
|
|
|
|
def __enter__(self):
|
|
self.old_context = multicore._CONTEXT
|
|
multicore._CONTEXT = None
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
multicore._CONTEXT = self.old_context
|
|
|
|
|
|
class Test__get_context(unittest.TestCase):
|
|
@unittest.skipUnless(not IS_SUPPORTING_CONTEXTS,
|
|
"Behaviour happens only in python <=3.3")
|
|
@mock.patch("imgaug.imgaug.warn")
|
|
@mock.patch("platform.version")
|
|
def test_mocked_nixos_python2(self, mock_version, mock_warn):
|
|
with clean_context():
|
|
mock_version.return_value = "NixOS"
|
|
_ctx = multicore._get_context()
|
|
assert mock_warn.call_count == 1
|
|
|
|
@unittest.skipUnless(IS_SUPPORTING_CONTEXTS,
|
|
"Behaviour is only supported in python 3.4+")
|
|
@mock.patch("platform.version")
|
|
@mock.patch("multiprocessing.get_context")
|
|
def test_mocked_nixos_python3(self, mock_gctx, mock_version):
|
|
with clean_context():
|
|
mock_version.return_value = "NixOS"
|
|
_ctx = multicore._get_context()
|
|
mock_gctx.assert_called_once_with("spawn")
|
|
|
|
@unittest.skipUnless(not IS_SUPPORTING_CONTEXTS,
|
|
"Behaviour happens only in python <=3.3")
|
|
@mock.patch("platform.version")
|
|
def test_mocked_no_nixos_python2(self, mock_version):
|
|
with clean_context():
|
|
mock_version.return_value = "Ubuntu"
|
|
ctx = multicore._get_context()
|
|
assert ctx is multiprocessing
|
|
|
|
@unittest.skipUnless(IS_SUPPORTING_CONTEXTS,
|
|
"Behaviour is only supported in python 3.4+")
|
|
@mock.patch("platform.system")
|
|
@mock.patch("multiprocessing.get_context")
|
|
@mock.patch("platform.version")
|
|
def test_mocked_no_nixos_python3(self, mock_version, mock_gctx, mock_system):
|
|
with clean_context():
|
|
mock_version.return_value = "Ubuntu"
|
|
mock_system.return_value = "Linux"
|
|
_ctx = multicore._get_context()
|
|
assert mock_gctx.call_count == 1
|
|
assert mock_gctx.call_args_list[0][0][0] is None
|
|
|
|
@unittest.skipUnless(IS_SUPPORTING_CONTEXTS,
|
|
"Behaviour is only supported in python 3.4+")
|
|
@mock.patch.object(sys, "version_info")
|
|
@mock.patch("platform.system")
|
|
@mock.patch("multiprocessing.get_context")
|
|
@mock.patch("platform.version")
|
|
def test_mocked_mac_and_37_cause_spawn(
|
|
self,
|
|
mock_version,
|
|
mock_gctx,
|
|
mock_system,
|
|
mock_vi
|
|
):
|
|
with clean_context():
|
|
def version_info(index):
|
|
if isinstance(index, slice):
|
|
return 3, 7
|
|
return 3 if index == 0 else 7
|
|
|
|
mock_vi.__getitem__.side_effect = version_info
|
|
|
|
mock_version.return_value = "foo"
|
|
mock_system.return_value = "Darwin"
|
|
_ctx = multicore._get_context()
|
|
mock_gctx.assert_called_once_with("spawn")
|
|
|
|
|
|
class TestPool(unittest.TestCase):
|
|
def setUp(self):
|
|
reseed()
|
|
|
|
def test___init___seed_out_of_bounds(self):
|
|
augseq = iaa.Identity()
|
|
with self.assertRaises(AssertionError) as context:
|
|
_ = multicore.Pool(augseq, seed=iarandom.SEED_MAX_VALUE + 100)
|
|
assert "Expected `seed` to be" in str(context.exception)
|
|
|
|
def test_property_pool(self):
|
|
mock_Pool = mock.MagicMock()
|
|
mock_Pool.return_value = mock_Pool
|
|
mock_Pool.close.return_value = None
|
|
mock_Pool.join.return_value = None
|
|
|
|
# We cannot just mock multiprocessing.Pool here, because of using
|
|
# a custom context. We would have to mock each possible context's
|
|
# Pool() method or overwrite here the Pool() method of the
|
|
# actually used context.
|
|
with mock.patch("multiprocessing.pool.Pool", mock_Pool):
|
|
augseq = iaa.Identity()
|
|
pool_config = multicore.Pool(
|
|
augseq, processes=1, maxtasksperchild=4, seed=123)
|
|
with pool_config as pool:
|
|
assert pool.processes == 1
|
|
assert pool._pool is None
|
|
assert mock_Pool.call_count == 1
|
|
assert mock_Pool.close.call_count == 1
|
|
assert mock_Pool.join.call_count == 1
|
|
# see
|
|
# https://github.com/
|
|
# python/cpython/blob/master/Lib/multiprocessing/context.py
|
|
# L119 (method Pool()) for an example of how Pool() is called
|
|
# internally.
|
|
assert mock_Pool.call_args[0][0] == 1 # processes
|
|
assert mock_Pool.call_args[0][1] is multicore._Pool_initialize_worker
|
|
assert mock_Pool.call_args[0][2] == (augseq, 123)
|
|
assert mock_Pool.call_args[0][3] == 4
|
|
|
|
def test_processes(self):
|
|
augseq = iaa.Identity()
|
|
mock_Pool = mock.MagicMock()
|
|
mock_cpu_count = mock.Mock()
|
|
|
|
# We cannot just mock multiprocessing.Pool here, because of using
|
|
# a custom context. We would have to mock each possible context's
|
|
# Pool() method or overwrite here the Pool() method of the
|
|
# actually used context.
|
|
patch_pool = mock.patch("multiprocessing.pool.Pool", mock_Pool)
|
|
|
|
# Multiprocessing seems to always access os.cpu_count to get the
|
|
# current count of cpu cores.
|
|
# See
|
|
# https://github.com/
|
|
# python/cpython/blob/master/Lib/multiprocessing/context.py
|
|
# L41.
|
|
fname = ("os.cpu_count" if IS_SUPPORTING_CONTEXTS
|
|
else "multiprocessing.cpu_count")
|
|
patch_cpu_count = mock.patch(fname, mock_cpu_count)
|
|
|
|
with patch_pool, patch_cpu_count:
|
|
# (cpu cores available, processes requested, processes started)
|
|
combos = [
|
|
(1, 1, 1),
|
|
(2, 1, 1),
|
|
(3, 1, 1),
|
|
(1, 2, 2),
|
|
(3, 2, 2),
|
|
(1, None, None),
|
|
(2, None, None),
|
|
(3, None, None),
|
|
(1, -1, 1),
|
|
(2, -1, 1),
|
|
(3, -1, 2),
|
|
(4, -2, 2)
|
|
]
|
|
|
|
for cores_available, processes_req, expected in combos:
|
|
with self.subTest(cpu_count_available=cores_available,
|
|
processes_requested=processes_req):
|
|
mock_cpu_count.return_value = cores_available
|
|
with multicore.Pool(augseq,
|
|
processes=processes_req) as _pool:
|
|
pass
|
|
|
|
if expected is None:
|
|
assert mock_Pool.call_args[0][0] is None
|
|
else:
|
|
assert mock_Pool.call_args[0][0] == expected
|
|
|
|
@mock.patch("multiprocessing.pool.Pool")
|
|
def test_cpu_count_does_not_exist(self, mock_pool):
|
|
def _side_effect():
|
|
raise NotImplementedError
|
|
|
|
old_method = multicore._get_context().cpu_count
|
|
mock_cpu_count = mock.Mock()
|
|
mock_cpu_count.side_effect = _side_effect
|
|
multicore._get_context().cpu_count = mock_cpu_count
|
|
|
|
augseq = iaa.Identity()
|
|
with warnings.catch_warnings(record=True) as caught_warnings:
|
|
warnings.simplefilter("always")
|
|
with multicore.Pool(augseq, processes=-1):
|
|
pass
|
|
|
|
assert mock_cpu_count.call_count == 1
|
|
assert mock_pool.call_count == 1
|
|
# 'processes' arg to Pool was expected to be set to None as cpu_count
|
|
# produced an error
|
|
assert mock_pool.call_args_list[0][0][0] is None
|
|
|
|
assert len(caught_warnings) == 1
|
|
assert (
|
|
"Could not find method multiprocessing.cpu_count(). "
|
|
in str(caught_warnings[-1].message))
|
|
|
|
multicore._get_context().cpu_count = old_method
|
|
|
|
@classmethod
|
|
def _test_map_batches_both(cls, call_async):
|
|
for clazz in [Batch, UnnormalizedBatch]:
|
|
augseq = iaa.Identity()
|
|
mock_Pool = mock.MagicMock()
|
|
mock_Pool.return_value = mock_Pool
|
|
mock_Pool.map.return_value = "X"
|
|
mock_Pool.map_async.return_value = "X"
|
|
with mock.patch("multiprocessing.pool.Pool", mock_Pool):
|
|
batches = [
|
|
clazz(images=[ia.data.quokka()]),
|
|
clazz(images=[ia.data.quokka()+1])
|
|
]
|
|
with multicore.Pool(augseq, processes=1) as pool:
|
|
if call_async:
|
|
_ = pool.map_batches_async(batches)
|
|
else:
|
|
_ = pool.map_batches(batches)
|
|
|
|
if call_async:
|
|
to_check = mock_Pool.map_async
|
|
else:
|
|
to_check = mock_Pool.map
|
|
|
|
assert to_check.call_count == 1
|
|
|
|
# args, arg 0
|
|
assert to_check.call_args[0][0] == multicore._Pool_starworker
|
|
|
|
# args, arg 1 (batches with ids), tuple 0,
|
|
# entry 0 in tuple (=> batch id)
|
|
assert to_check.call_args[0][1][0][0] == 0
|
|
|
|
# args, arg 1 (batches with ids), tuple 0,
|
|
# entry 1 in tuple (=> batch)
|
|
assert np.array_equal(
|
|
to_check.call_args[0][1][0][1].images_unaug,
|
|
batches[0].images_unaug)
|
|
|
|
# args, arg 1 (batches with ids), tuple 1,
|
|
# entry 0 in tuple (=> batch id)
|
|
assert to_check.call_args[0][1][1][0] == 1
|
|
|
|
# args, arg 1 (batches with ids), tuple 1,
|
|
# entry 1 in tuple (=> batch)
|
|
assert np.array_equal(
|
|
to_check.call_args[0][1][1][1].images_unaug,
|
|
batches[1].images_unaug)
|
|
|
|
def test_map_batches(self):
|
|
self._test_map_batches_both(call_async=False)
|
|
|
|
def test_map_batches_async(self):
|
|
self._test_map_batches_both(call_async=True)
|
|
|
|
@classmethod
|
|
def _test_imap_batches_both(cls, call_unordered):
|
|
for clazz in [Batch, UnnormalizedBatch]:
|
|
batches = [clazz(images=[ia.data.quokka()]),
|
|
clazz(images=[ia.data.quokka()+1])]
|
|
|
|
def _generate_batches():
|
|
for batch in batches:
|
|
yield batch
|
|
|
|
augseq = iaa.Identity()
|
|
mock_Pool = mock.MagicMock()
|
|
mock_Pool.return_value = mock_Pool
|
|
mock_Pool.imap.return_value = batches
|
|
mock_Pool.imap_unordered.return_value = batches
|
|
with mock.patch("multiprocessing.pool.Pool", mock_Pool):
|
|
with multicore.Pool(augseq, processes=1) as pool:
|
|
gen = _generate_batches()
|
|
if call_unordered:
|
|
_ = list(pool.imap_batches_unordered(gen))
|
|
else:
|
|
_ = list(pool.imap_batches(gen))
|
|
|
|
if call_unordered:
|
|
to_check = mock_Pool.imap_unordered
|
|
else:
|
|
to_check = mock_Pool.imap
|
|
|
|
assert to_check.call_count == 1
|
|
|
|
assert to_check.call_args[0][0] == multicore._Pool_starworker
|
|
|
|
# convert generator to list, make it subscriptable
|
|
arg_batches = list(to_check.call_args[0][1])
|
|
|
|
# args, arg 1 (batches with ids), tuple 0,
|
|
# entry 0 in tuple (=> batch id)
|
|
assert arg_batches[0][0] == 0
|
|
|
|
# tuple 0, entry 1 in tuple (=> batch)
|
|
assert np.array_equal(
|
|
arg_batches[0][1].images_unaug,
|
|
batches[0].images_unaug)
|
|
|
|
# tuple 1, entry 0 in tuple (=> batch id)
|
|
assert arg_batches[1][0] == 1
|
|
|
|
# tuple 1, entry 1 in tuple (=> batch)
|
|
assert np.array_equal(
|
|
arg_batches[1][1].images_unaug,
|
|
batches[1].images_unaug)
|
|
|
|
@classmethod
|
|
def _test_imap_batches_both_output_buffer_size(cls, call_unordered,
|
|
timeout=0.075):
|
|
batches = [
|
|
ia.Batch(images=[np.full((1, 1), i, dtype=np.uint8)])
|
|
for i in range(8)]
|
|
|
|
def _generate_batches(times):
|
|
for batch in batches:
|
|
yield batch
|
|
times.append(time.time())
|
|
|
|
def callfunc(pool, gen, output_buffer_size):
|
|
func = (
|
|
pool.imap_batches_unordered
|
|
if call_unordered
|
|
else pool.imap_batches
|
|
)
|
|
|
|
for v in func(gen, output_buffer_size=output_buffer_size):
|
|
yield v
|
|
|
|
def contains_all_ids(inputs):
|
|
arrs = np.uint8([batch.images_aug for batch in inputs])
|
|
ids_uq = np.unique(arrs)
|
|
return (
|
|
len(ids_uq) == len(batches)
|
|
and np.all(0 <= ids_uq)
|
|
and np.all(ids_uq < len(batches))
|
|
)
|
|
|
|
augseq = iaa.Identity()
|
|
with multicore.Pool(augseq, processes=1) as pool:
|
|
# no output buffer limit, there should be no noteworthy lag
|
|
# for any batch requested from _generate_batches()
|
|
times = []
|
|
gen = callfunc(pool, _generate_batches(times), None)
|
|
result = next(gen)
|
|
time.sleep(timeout)
|
|
result = [result] + list(gen)
|
|
times = np.float64(times)
|
|
times_diffs = times[1:] - times[0:-1]
|
|
assert np.all(times_diffs < timeout * 1.01)
|
|
assert contains_all_ids(result)
|
|
|
|
# with output buffer limit, but set to the number of batches,
|
|
# i.e. should again not lead to any lag
|
|
times = []
|
|
gen = callfunc(pool, _generate_batches(times), len(batches))
|
|
result = next(gen)
|
|
time.sleep(timeout)
|
|
result = [result] + list(gen)
|
|
times = np.float64(times)
|
|
times_diffs = times[1:] - times[0:-1]
|
|
assert np.all(times_diffs < timeout * 1.01)
|
|
assert contains_all_ids(result)
|
|
|
|
# With output buffer limit of #batches/2 (=4), followed by a
|
|
# timeout after starting the loading process. This should quickly
|
|
# load batches until the buffer is full, then wait until the
|
|
# batches are requested from the buffer (i.e. after the timeout
|
|
# ended) and then proceed to produce batches at the speed at which
|
|
# they are requested. This should lead to a measureable lag between
|
|
# batch 4 and 5 (matching the timeout).
|
|
times = []
|
|
gen = callfunc(pool, _generate_batches(times), 4)
|
|
result = next(gen)
|
|
time.sleep(timeout)
|
|
result = [result] + list(gen)
|
|
times = np.float64(times)
|
|
times_diffs = times[1:] - times[0:-1]
|
|
# use -1 here because we have N-1 times for N batches as
|
|
# diffs denote diffs between Nth and N+1th batch
|
|
assert np.all(times_diffs[0:4-1] < timeout * 1.01)
|
|
assert np.all(times_diffs[4-1:4-1+1] >= timeout * 0.99)
|
|
assert np.all(times_diffs[4-1+1:] < timeout * 1.01)
|
|
assert contains_all_ids(result)
|
|
|
|
def test_imap_batches(self):
|
|
self._test_imap_batches_both(call_unordered=False)
|
|
|
|
def test_imap_batches_unordered(self):
|
|
self._test_imap_batches_both(call_unordered=True)
|
|
|
|
def test_imap_batches_output_buffer_size(self):
|
|
self._test_imap_batches_both_output_buffer_size(call_unordered=False)
|
|
|
|
def test_imap_batches_unordered_output_buffer_size(self):
|
|
self._test_imap_batches_both_output_buffer_size(call_unordered=True)
|
|
|
|
@classmethod
|
|
def _assert_each_augmentation_not_more_than_once(cls, batches_aug):
|
|
sum_to_vecs = defaultdict(list)
|
|
for batch in batches_aug:
|
|
assert not np.array_equal(batch.images_aug[0], batch.images_aug[1])
|
|
|
|
vec = batch.images_aug.flatten()
|
|
vecsum = int(np.sum(vec))
|
|
if vecsum in sum_to_vecs:
|
|
for other_vec in sum_to_vecs[vecsum]:
|
|
assert not np.array_equal(vec, other_vec)
|
|
else:
|
|
sum_to_vecs[vecsum].append(vec)
|
|
|
|
def test_augmentations_with_seed_match(self):
|
|
nb_batches = 60
|
|
augseq = iaa.AddElementwise((0, 255))
|
|
image = np.zeros((10, 10, 1), dtype=np.uint8)
|
|
batch = ia.Batch(images=np.uint8([image, image]))
|
|
batches = [batch.deepcopy() for _ in sm.xrange(nb_batches)]
|
|
|
|
# seed=1
|
|
with multicore.Pool(augseq, processes=2, maxtasksperchild=30,
|
|
seed=1) as pool:
|
|
batches_aug1 = pool.map_batches(batches, chunksize=2)
|
|
|
|
# seed=1
|
|
with multicore.Pool(augseq, processes=2, seed=1) as pool:
|
|
batches_aug2 = pool.map_batches(batches, chunksize=1)
|
|
# seed=2
|
|
with multicore.Pool(augseq, processes=2, seed=2) as pool:
|
|
batches_aug3 = pool.map_batches(batches, chunksize=1)
|
|
|
|
assert len(batches_aug1) == nb_batches
|
|
assert len(batches_aug2) == nb_batches
|
|
assert len(batches_aug3) == nb_batches
|
|
|
|
for b1, b2, b3 in zip(batches_aug1, batches_aug2, batches_aug3):
|
|
# images were augmented
|
|
assert not np.array_equal(b1.images_unaug, b1.images_aug)
|
|
assert not np.array_equal(b2.images_unaug, b2.images_aug)
|
|
assert not np.array_equal(b3.images_unaug, b3.images_aug)
|
|
|
|
# original images still the same
|
|
assert np.array_equal(b1.images_unaug, batch.images_unaug)
|
|
assert np.array_equal(b2.images_unaug, batch.images_unaug)
|
|
assert np.array_equal(b3.images_unaug, batch.images_unaug)
|
|
|
|
# augmentations for same seed are the same
|
|
assert np.array_equal(b1.images_aug, b2.images_aug)
|
|
|
|
# augmentations for different seeds are different
|
|
assert not np.array_equal(b1.images_aug, b3.images_aug)
|
|
|
|
# make sure that batches for the two pools with same seed did not
|
|
# repeat within results (only between the results of the two pools)
|
|
for batches_aug in [batches_aug1, batches_aug2, batches_aug3]:
|
|
self._assert_each_augmentation_not_more_than_once(batches_aug)
|
|
|
|
def test_augmentations_with_seed_match_for_images_and_keypoints(self):
|
|
augseq = iaa.AddElementwise((0, 255))
|
|
image = np.zeros((10, 10, 1), dtype=np.uint8)
|
|
# keypoints here will not be changed by augseq, but they will induce
|
|
# deterministic mode to start in augment_batches() as each batch
|
|
# contains images AND keypoints
|
|
kps = ia.KeypointsOnImage([ia.Keypoint(x=2, y=0)], shape=(10, 10, 1))
|
|
batch = ia.Batch(images=np.uint8([image, image]), keypoints=[kps, kps])
|
|
batches = [batch.deepcopy() for _ in sm.xrange(60)]
|
|
|
|
# seed=1
|
|
with multicore.Pool(augseq, processes=2, maxtasksperchild=30,
|
|
seed=1) as pool:
|
|
batches_aug1 = pool.map_batches(batches, chunksize=2)
|
|
# seed=1
|
|
with multicore.Pool(augseq, processes=2, seed=1) as pool:
|
|
batches_aug2 = pool.map_batches(batches, chunksize=1)
|
|
# seed=2
|
|
with multicore.Pool(augseq, processes=2, seed=2) as pool:
|
|
batches_aug3 = pool.map_batches(batches, chunksize=1)
|
|
|
|
assert len(batches_aug1) == 60
|
|
assert len(batches_aug2) == 60
|
|
assert len(batches_aug3) == 60
|
|
|
|
for batches_aug in [batches_aug1, batches_aug2, batches_aug3]:
|
|
for batch in batches_aug:
|
|
for keypoints_aug in batch.keypoints_aug:
|
|
assert keypoints_aug.keypoints[0].x == 2
|
|
assert keypoints_aug.keypoints[0].y == 0
|
|
|
|
for b1, b2, b3 in zip(batches_aug1, batches_aug2, batches_aug3):
|
|
# images were augmented
|
|
assert not np.array_equal(b1.images_unaug, b1.images_aug)
|
|
assert not np.array_equal(b2.images_unaug, b2.images_aug)
|
|
assert not np.array_equal(b3.images_unaug, b3.images_aug)
|
|
|
|
# original images still the same
|
|
assert np.array_equal(b1.images_unaug, batch.images_unaug)
|
|
assert np.array_equal(b2.images_unaug, batch.images_unaug)
|
|
assert np.array_equal(b3.images_unaug, batch.images_unaug)
|
|
|
|
# augmentations for same seed are the same
|
|
assert np.array_equal(b1.images_aug, b2.images_aug)
|
|
|
|
# augmentations for different seeds are different
|
|
assert not np.array_equal(b1.images_aug, b3.images_aug)
|
|
|
|
# make sure that batches for the two pools with same seed did not
|
|
# repeat within results (only between the results of the two pools)
|
|
for batches_aug in [batches_aug1, batches_aug2, batches_aug3]:
|
|
self._assert_each_augmentation_not_more_than_once(batches_aug)
|
|
|
|
def test_augmentations_without_seed_differ(self):
|
|
augseq = iaa.AddElementwise((0, 255))
|
|
image = np.zeros((10, 10, 1), dtype=np.uint8)
|
|
batch = ia.Batch(images=np.uint8([image, image]))
|
|
batches = [batch.deepcopy() for _ in sm.xrange(20)]
|
|
|
|
with multicore.Pool(augseq, processes=2, maxtasksperchild=5) as pool:
|
|
batches_aug = pool.map_batches(batches, chunksize=2)
|
|
with multicore.Pool(augseq, processes=2) as pool:
|
|
batches_aug.extend(pool.map_batches(batches, chunksize=1))
|
|
|
|
assert len(batches_aug) == 2*20
|
|
|
|
self._assert_each_augmentation_not_more_than_once(batches_aug)
|
|
|
|
def test_augmentations_without_seed_differ_for_images_and_keypoints(self):
|
|
augseq = iaa.AddElementwise((0, 255))
|
|
image = np.zeros((10, 10, 1), dtype=np.uint8)
|
|
# keypoints here will not be changed by augseq, but they will
|
|
# induce deterministic mode to start in augment_batches() as each
|
|
# batch contains images AND keypoints
|
|
kps = ia.KeypointsOnImage([ia.Keypoint(x=2, y=0)], shape=(10, 10, 1))
|
|
batch = ia.Batch(images=np.uint8([image, image]), keypoints=[kps, kps])
|
|
batches = [batch.deepcopy() for _ in sm.xrange(20)]
|
|
|
|
with multicore.Pool(augseq, processes=2, maxtasksperchild=5) as pool:
|
|
batches_aug = pool.map_batches(batches, chunksize=2)
|
|
with multicore.Pool(augseq, processes=2) as pool:
|
|
batches_aug.extend(pool.map_batches(batches, chunksize=1))
|
|
|
|
assert len(batches_aug) == 2*20
|
|
|
|
for batch in batches_aug:
|
|
for keypoints_aug in batch.keypoints_aug:
|
|
assert keypoints_aug.keypoints[0].x == 2
|
|
assert keypoints_aug.keypoints[0].y == 0
|
|
|
|
self._assert_each_augmentation_not_more_than_once(batches_aug)
|
|
|
|
def test_inputs_not_lost(self):
|
|
"""Test to make sure that inputs (e.g. images) are never lost."""
|
|
def _assert_contains_all_ids(batches_aug):
|
|
# batch.images_unaug
|
|
ids = set()
|
|
for batch_aug in batches_aug:
|
|
ids.add(int(batch_aug.images_unaug.flat[0]))
|
|
ids.add(int(batch_aug.images_unaug.flat[1]))
|
|
for idx in sm.xrange(2*100):
|
|
assert idx in ids
|
|
assert len(ids) == 200
|
|
|
|
# batch.images_aug
|
|
ids = set()
|
|
for batch_aug in batches_aug:
|
|
ids.add(int(batch_aug.images_aug.flat[0]))
|
|
ids.add(int(batch_aug.images_aug.flat[1]))
|
|
for idx in sm.xrange(2*100):
|
|
assert idx in ids
|
|
assert len(ids) == 200
|
|
|
|
augseq = iaa.Identity()
|
|
image = np.zeros((1, 1, 1), dtype=np.uint8)
|
|
# creates batches containing images with ids from 0 to 199 (one pair
|
|
# of consecutive ids per batch)
|
|
batches = [
|
|
ia.Batch(images=np.uint8([image + b_idx*2, image + b_idx*2+1]))
|
|
for b_idx
|
|
in sm.xrange(100)]
|
|
|
|
with multicore.Pool(augseq, processes=2, maxtasksperchild=25) as pool:
|
|
batches_aug = pool.map_batches(batches)
|
|
_assert_contains_all_ids(batches_aug)
|
|
|
|
with multicore.Pool(augseq, processes=2, maxtasksperchild=25,
|
|
seed=1) as pool:
|
|
batches_aug = pool.map_batches(batches)
|
|
_assert_contains_all_ids(batches_aug)
|
|
|
|
with multicore.Pool(augseq, processes=3, seed=2) as pool:
|
|
batches_aug = pool.map_batches(batches)
|
|
_assert_contains_all_ids(batches_aug)
|
|
|
|
with multicore.Pool(augseq, processes=2, seed=None) as pool:
|
|
batches_aug = pool.map_batches(batches)
|
|
_assert_contains_all_ids(batches_aug)
|
|
|
|
batches_aug = pool.map_batches(batches)
|
|
_assert_contains_all_ids(batches_aug)
|
|
|
|
def test_close(self):
|
|
augseq = iaa.Identity()
|
|
with multicore.Pool(augseq, processes=2) as pool:
|
|
pool.close()
|
|
|
|
def test_terminate(self):
|
|
augseq = iaa.Identity()
|
|
with multicore.Pool(augseq, processes=2) as pool:
|
|
pool.terminate()
|
|
|
|
def test_join(self):
|
|
augseq = iaa.Identity()
|
|
with multicore.Pool(augseq, processes=2) as pool:
|
|
pool.close()
|
|
pool.join()
|
|
|
|
@mock.patch("multiprocessing.pool.Pool")
|
|
def test_join_via_mock(self, mock_pool):
|
|
# According to codecov, the join() does not get beyond its initial
|
|
# if statement in the test_join() test, even though it should be.
|
|
# Might be a simple travis multicore problem?
|
|
# It is tested here again via some mocking.
|
|
mock_pool.return_value = mock_pool
|
|
mock_pool.join.return_value = True
|
|
with multicore.Pool(iaa.Identity(), processes=2) as pool:
|
|
pool.join()
|
|
|
|
# Make sure that __exit__ does not call close(), which would then
|
|
# call join() again and we would get a call_count of 2
|
|
pool._pool = None
|
|
|
|
assert mock_pool.join.call_count == 1
|
|
|
|
|
|
# This should already be part of the Pool tests, but according to codecov
|
|
# it is not tested. Likely some travis error related to running multiple
|
|
# python processes.
|
|
class Test_Pool_initialize_worker(unittest.TestCase):
|
|
def tearDown(self):
|
|
# without this, other tests can break as e.g. the functions in
|
|
# multicore assert that _WORKER_AUGSEQ is None
|
|
multicore.Pool._WORKER_AUGSEQ = None
|
|
multicore.Pool._WORKER_SEED_START = None
|
|
|
|
@mock.patch("imgaug.multicore.Pool")
|
|
def test_with_seed_start(self, mock_ia_pool):
|
|
augseq = mock.MagicMock()
|
|
multicore._Pool_initialize_worker(augseq, 1)
|
|
assert mock_ia_pool._WORKER_SEED_START == 1
|
|
assert mock_ia_pool._WORKER_AUGSEQ is augseq
|
|
assert augseq.localize_random_state_.call_count == 1
|
|
|
|
@mock.patch.object(sys, 'version_info')
|
|
@mock.patch("time.time_ns", create=True) # doesnt exist in <=3.6
|
|
@mock.patch("imgaug.random.seed")
|
|
@mock.patch("multiprocessing.current_process")
|
|
def test_without_seed_start_simulate_py37_or_higher(self,
|
|
mock_cp,
|
|
mock_ia_seed,
|
|
mock_time_ns,
|
|
mock_vi):
|
|
def version_info(index):
|
|
return 3 if index == 0 else 7
|
|
|
|
mock_vi.__getitem__.side_effect = version_info
|
|
mock_time_ns.return_value = 1
|
|
mock_cp.return_value = mock.MagicMock()
|
|
mock_cp.return_value.name = "foo"
|
|
augseq = mock.MagicMock()
|
|
|
|
multicore._Pool_initialize_worker(augseq, None)
|
|
|
|
assert mock_time_ns.call_count == 1
|
|
assert mock_ia_seed.call_count == 1
|
|
assert augseq.seed_.call_count == 1
|
|
|
|
seed_global = mock_ia_seed.call_args_list[0][0][0]
|
|
seed_local = augseq.seed_.call_args_list[0][0][0]
|
|
assert seed_global != seed_local
|
|
|
|
@mock.patch.object(sys, 'version_info')
|
|
@mock.patch("time.time")
|
|
@mock.patch("imgaug.random.seed")
|
|
@mock.patch("multiprocessing.current_process")
|
|
def test_without_seed_start_simulate_py36_or_lower(self,
|
|
mock_cp,
|
|
mock_ia_seed,
|
|
mock_time,
|
|
mock_vi):
|
|
def version_info(index):
|
|
return 3 if index == 0 else 6
|
|
|
|
mock_vi.__getitem__.side_effect = version_info
|
|
mock_time.return_value = 1
|
|
mock_cp.return_value = mock.MagicMock()
|
|
mock_cp.return_value.name = "foo"
|
|
augseq = mock.MagicMock()
|
|
|
|
multicore._Pool_initialize_worker(augseq, None)
|
|
|
|
assert mock_time.call_count == 1
|
|
assert mock_ia_seed.call_count == 1
|
|
assert augseq.seed_.call_count == 1
|
|
|
|
seed_global = mock_ia_seed.call_args_list[0][0][0]
|
|
seed_local = augseq.seed_.call_args_list[0][0][0]
|
|
assert seed_global != seed_local
|
|
|
|
@mock.patch("imgaug.random.seed")
|
|
def test_without_seed_start(self, mock_ia_seed):
|
|
augseq = mock.MagicMock()
|
|
|
|
multicore._Pool_initialize_worker(augseq, None)
|
|
time.sleep(0.01)
|
|
multicore._Pool_initialize_worker(augseq, None)
|
|
|
|
seed_global_call_1 = mock_ia_seed.call_args_list[0][0][0]
|
|
seed_local_call_1 = augseq.seed_.call_args_list[0][0][0]
|
|
seed_global_call_2 = mock_ia_seed.call_args_list[0][0][0]
|
|
seed_local_call_2 = augseq.seed_.call_args_list[0][0][0]
|
|
assert (
|
|
seed_global_call_1
|
|
!= seed_local_call_1
|
|
!= seed_global_call_2
|
|
!= seed_local_call_2
|
|
), "Got seeds: %d, %d, %d, %d" % (
|
|
seed_global_call_1, seed_local_call_1,
|
|
seed_global_call_2, seed_local_call_2)
|
|
assert mock_ia_seed.call_count == 2
|
|
assert augseq.seed_.call_count == 2
|
|
|
|
|
|
# This should already be part of the Pool tests, but according to codecov
|
|
# it is not tested. Likely some travis error related to running multiple
|
|
# python processes.
|
|
class Test_Pool_worker(unittest.TestCase):
|
|
def tearDown(self):
|
|
# without this, other tests can break as e.g. the functions in
|
|
# multicore assert that _WORKER_AUGSEQ is None
|
|
multicore.Pool._WORKER_AUGSEQ = None
|
|
multicore.Pool._WORKER_SEED_START = None
|
|
|
|
def test_without_seed_start(self):
|
|
augseq = mock.MagicMock()
|
|
augseq.augment_batch_.return_value = "augmented_batch_"
|
|
image = np.zeros((1, 1, 3), dtype=np.uint8)
|
|
batch = UnnormalizedBatch(images=[image])
|
|
|
|
multicore.Pool._WORKER_AUGSEQ = augseq
|
|
result = multicore._Pool_worker(1, batch)
|
|
|
|
assert result == "augmented_batch_"
|
|
assert augseq.augment_batch_.call_count == 1
|
|
augseq.augment_batch_.assert_called_once_with(batch)
|
|
|
|
@mock.patch("imgaug.random.seed")
|
|
def test_with_seed_start(self, mock_ia_seed):
|
|
augseq = mock.MagicMock()
|
|
augseq.augment_batch_.return_value = "augmented_batch_"
|
|
image = np.zeros((1, 1, 3), dtype=np.uint8)
|
|
batch = UnnormalizedBatch(images=[image])
|
|
batch_idx = 1
|
|
seed_start = 10
|
|
|
|
multicore.Pool._WORKER_AUGSEQ = augseq
|
|
multicore.Pool._WORKER_SEED_START = seed_start
|
|
result = multicore._Pool_worker(batch_idx, batch)
|
|
|
|
# expected seeds used
|
|
seed = seed_start + batch_idx
|
|
seed_global_expected = (
|
|
iarandom.SEED_MIN_VALUE
|
|
+ (seed - 10**9)
|
|
% (iarandom.SEED_MAX_VALUE - iarandom.SEED_MIN_VALUE)
|
|
)
|
|
seed_local_expected = (
|
|
iarandom.SEED_MIN_VALUE
|
|
+ seed
|
|
% (iarandom.SEED_MAX_VALUE - iarandom.SEED_MIN_VALUE)
|
|
)
|
|
|
|
assert result == "augmented_batch_"
|
|
assert augseq.augment_batch_.call_count == 1
|
|
augseq.augment_batch_.assert_called_once_with(batch)
|
|
mock_ia_seed.assert_called_once_with(seed_global_expected)
|
|
augseq.seed_.assert_called_once_with(seed_local_expected)
|
|
|
|
|
|
# This should already be part of the Pool tests, but according to codecov
|
|
# it is not tested. Likely some travis error related to running multiple
|
|
# python processes.
|
|
class Test_Pool_starworker(unittest.TestCase):
|
|
def tearDown(self):
|
|
# without this, other tests can break as e.g. the functions in
|
|
# multicore assert that _WORKER_AUGSEQ is None
|
|
multicore.Pool._WORKER_AUGSEQ = None
|
|
multicore.Pool._WORKER_SEED_START = None
|
|
|
|
@mock.patch("imgaug.multicore._Pool_worker")
|
|
def test_simple_call(self, mock_worker):
|
|
image = np.zeros((1, 1, 3), dtype=np.uint8)
|
|
batch = UnnormalizedBatch(images=[image])
|
|
batch_idx = 1
|
|
mock_worker.return_value = "returned_batch"
|
|
|
|
result = multicore._Pool_starworker((batch_idx, batch))
|
|
|
|
assert result == "returned_batch"
|
|
mock_worker.assert_called_once_with(batch_idx, batch)
|
|
|
|
|
|
# ---------
|
|
# loading function used in TestBatchLoader.test_basic_functionality()
|
|
# it is outside of the test as putting it inside of it caused issues
|
|
# with spawn mode not being able to pickle this method, see issue #414.
|
|
def _batch_loader_load_func():
|
|
for _ in sm.xrange(20):
|
|
yield ia.Batch(images=np.zeros((2, 4, 4, 3), dtype=np.uint8))
|
|
# ---------
|
|
|
|
|
|
# Note that BatchLoader is deprecated
|
|
class TestBatchLoader(unittest.TestCase):
|
|
def setUp(self):
|
|
reseed()
|
|
|
|
def test_basic_functionality(self):
|
|
warnings.simplefilter("always")
|
|
with warnings.catch_warnings(record=True) as caught_warnings:
|
|
for nb_workers in [1, 2]:
|
|
# repeat these tests many times to catch rarer race conditions
|
|
for _ in sm.xrange(5):
|
|
loader = multicore.BatchLoader(
|
|
_batch_loader_load_func, queue_size=2,
|
|
nb_workers=nb_workers, threaded=True)
|
|
loaded = []
|
|
counter = 0
|
|
while ((not loader.all_finished()
|
|
or not loader.queue.empty())
|
|
and counter < 1000):
|
|
try:
|
|
batch = loader.queue.get(timeout=0.001)
|
|
loaded.append(batch)
|
|
except:
|
|
pass
|
|
counter += 1
|
|
assert len(loaded) == 20*nb_workers, \
|
|
"Expected %d to be loaded by threads, got %d for %d " \
|
|
"workers at counter %d." % (
|
|
20*nb_workers, len(loaded), nb_workers, counter
|
|
)
|
|
|
|
loader = multicore.BatchLoader(
|
|
_batch_loader_load_func, queue_size=200,
|
|
nb_workers=nb_workers, threaded=True)
|
|
loader.terminate()
|
|
assert loader.all_finished()
|
|
|
|
loader = multicore.BatchLoader(
|
|
_batch_loader_load_func, queue_size=2,
|
|
nb_workers=nb_workers, threaded=False)
|
|
loaded = []
|
|
counter = 0
|
|
while ((not loader.all_finished()
|
|
or not loader.queue.empty())
|
|
and counter < 1000):
|
|
try:
|
|
batch = loader.queue.get(timeout=0.001)
|
|
loaded.append(batch)
|
|
except:
|
|
pass
|
|
counter += 1
|
|
assert len(loaded) == 20*nb_workers, \
|
|
"Expected %d to be loaded by background processes, " \
|
|
"got %d for %d workers at counter %d." % (
|
|
20*nb_workers, len(loaded), nb_workers, counter
|
|
)
|
|
|
|
loader = multicore.BatchLoader(
|
|
_batch_loader_load_func, queue_size=200,
|
|
nb_workers=nb_workers, threaded=False)
|
|
loader.terminate()
|
|
assert loader.all_finished()
|
|
|
|
assert len(caught_warnings) > 0
|
|
for warning in caught_warnings:
|
|
assert "is deprecated" in str(warning.message)
|
|
|
|
|
|
# Note that BackgroundAugmenter is deprecated
|
|
class TestBackgroundAugmenter(unittest.TestCase):
|
|
def setUp(self):
|
|
reseed()
|
|
|
|
def test_augment_images_worker(self):
|
|
warnings.simplefilter("always")
|
|
with warnings.catch_warnings(record=True) as caught_warnings:
|
|
def gen():
|
|
yield ia.Batch(images=np.zeros((1, 4, 4, 3), dtype=np.uint8))
|
|
bl = multicore.BatchLoader(gen(), queue_size=2)
|
|
bgaug = multicore.BackgroundAugmenter(bl, iaa.Identity(),
|
|
queue_size=1, nb_workers=1)
|
|
|
|
queue_source = multiprocessing.Queue(2)
|
|
queue_target = multiprocessing.Queue(2)
|
|
queue_source.put(
|
|
pickle.dumps(
|
|
ia.Batch(images=np.zeros((1, 4, 8, 3), dtype=np.uint8)),
|
|
protocol=-1
|
|
)
|
|
)
|
|
queue_source.put(pickle.dumps(None, protocol=-1))
|
|
bgaug._augment_images_worker(iaa.Add(1), queue_source,
|
|
queue_target, 1)
|
|
|
|
batch_aug = pickle.loads(queue_target.get())
|
|
assert isinstance(batch_aug, ia.Batch)
|
|
assert batch_aug.images_unaug is not None
|
|
assert batch_aug.images_unaug.dtype == np.uint8
|
|
assert batch_aug.images_unaug.shape == (1, 4, 8, 3)
|
|
assert np.array_equal(
|
|
batch_aug.images_unaug,
|
|
np.zeros((1, 4, 8, 3), dtype=np.uint8))
|
|
assert batch_aug.images_aug is not None
|
|
assert batch_aug.images_aug.dtype == np.uint8
|
|
assert batch_aug.images_aug.shape == (1, 4, 8, 3)
|
|
assert np.array_equal(
|
|
batch_aug.images_aug,
|
|
np.zeros((1, 4, 8, 3), dtype=np.uint8) + 1)
|
|
|
|
finished_signal = pickle.loads(queue_target.get())
|
|
assert finished_signal is None
|
|
|
|
source_finished_signal = pickle.loads(queue_source.get())
|
|
assert source_finished_signal is None
|
|
|
|
assert queue_source.empty()
|
|
assert queue_target.empty()
|
|
|
|
queue_source.close()
|
|
queue_target.close()
|
|
queue_source.join_thread()
|
|
queue_target.join_thread()
|
|
bl.terminate()
|
|
bgaug.terminate()
|
|
|
|
assert len(caught_warnings) > 0
|
|
for warning in caught_warnings:
|
|
assert "is deprecated" in str(warning.message)
|