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

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)