def test_BackgroundAugmenter__augment_images_worker():
reseed()
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.Noop(),
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 ("BatchLoader is deprecated" in str(warning.message)
or "BackgroundAugmenter is deprecated" in str(warning.message))
def _augment_small_2():
batch_loader = multicore.BatchLoader(load_images(n_batches=2), queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq_i)
i = 0
while True:
batch = bg_augmenter.get_batch()
if batch is None:
break
i += 1
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)
def _augment_small_4():
batch_loader = multicore.BatchLoader(load_images(n_batches=10),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
batch = bg_augmenter.get_batch()
def main():
augseq = iaa.Sequential(
[iaa.Fliplr(0.5),
iaa.CoarseDropout(p=0.1, size_percent=0.1)])
def func_images(images, random_state, parents, hooks):
time.sleep(0.2)
return images
def func_heatmaps(heatmaps, random_state, parents, hooks):
return heatmaps
def func_keypoints(keypoints_on_images, random_state, parents, hooks):
return keypoints_on_images
augseq_slow = iaa.Sequential([
iaa.Fliplr(0.5),
iaa.Lambda(func_images=func_images,
func_heatmaps=func_heatmaps,
func_keypoints=func_keypoints)
])
print("------------------")
print("augseq.augment_batches(batches, background=True)")
print("------------------")
batches = list(load_images())
batches_aug = augseq.augment_batches(batches, background=True)
images_aug = []
keypoints_aug = []
for batch_aug in batches_aug:
images_aug.append(batch_aug.images_aug)
keypoints_aug.append(batch_aug.keypoints_aug)
ia.imshow(draw_grid(images_aug, keypoints_aug))
print("------------------")
print("augseq.augment_batches(batches, background=True) -> only images")
print("------------------")
batches = list(load_images())
batches = [batch.images_unaug for batch in batches]
batches_aug = augseq.augment_batches(batches, background=True)
images_aug = []
keypoints_aug = None
for batch_aug in batches_aug:
images_aug.append(batch_aug)
ia.imshow(draw_grid(images_aug, keypoints_aug))
print("------------------")
print("BackgroundAugmenter")
print("------------------")
batch_loader = multicore.BatchLoader(load_images)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq)
images_aug = []
keypoints_aug = []
while True:
print("Next batch...")
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
images_aug.append(batch.images_aug)
keypoints_aug.append(batch.keypoints_aug)
ia.imshow(draw_grid(images_aug, keypoints_aug))
print("------------------")
print("BackgroundAugmenter with generator in BL")
print("------------------")
batch_loader = multicore.BatchLoader(load_images())
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq)
images_aug = []
keypoints_aug = []
while True:
print("Next batch...")
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
images_aug.append(batch.images_aug)
keypoints_aug.append(batch.keypoints_aug)
ia.imshow(draw_grid(images_aug, keypoints_aug))
print("------------------")
print("Long running BackgroundAugmenter at BL-queue_size=12")
print("------------------")
batch_loader = multicore.BatchLoader(load_images(n_batches=1000),
queue_size=12)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq)
i = 0
while True:
if i % 100 == 0:
print("batch=%d..." % (i, ))
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
i += 1
print("------------------")
print("Long running BackgroundAugmenter at BL-queue_size=2")
print("------------------")
batch_loader = multicore.BatchLoader(load_images(n_batches=1000),
queue_size=2)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq)
i = 0
while True:
if i % 100 == 0:
print("batch=%d..." % (i, ))
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
i += 1
print("------------------")
print("Long running BackgroundAugmenter (slow loading)")
print("------------------")
batch_loader = multicore.BatchLoader(load_images(n_batches=100, sleep=0.2))
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq)
i = 0
while True:
if i % 10 == 0:
print("batch=%d..." % (i, ))
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
i += 1
print("------------------")
print("Long running BackgroundAugmenter (slow aug) at BL-queue_size=12")
print("------------------")
batch_loader = multicore.BatchLoader(load_images(n_batches=100),
queue_size=12)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq_slow)
i = 0
while True:
if i % 10 == 0:
print("batch=%d..." % (i, ))
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
i += 1
print("------------------")
print("Long running BackgroundAugmenter (slow aug) at BL-queue_size=2")
print("------------------")
batch_loader = multicore.BatchLoader(load_images(n_batches=100),
queue_size=2)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq_slow)
i = 0
while True:
if i % 10 == 0:
print("batch=%d..." % (i, ))
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished.")
break
i += 1
for augseq_i in [augseq, augseq_slow]:
print("------------------")
print("Many very small runs (batches=1)")
print("------------------")
for i in range(100):
batch_loader = multicore.BatchLoader(load_images(n_batches=1),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
while True:
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished (%d/%d)." % (i + 1, 100))
break
print("------------------")
print("Many very small runs (batches=2)")
print("------------------")
for i in range(100):
batch_loader = multicore.BatchLoader(load_images(n_batches=2),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
while True:
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished (%d/%d)." % (i + 1, 100))
break
print("------------------")
print("Many very small runs, separate function (batches=1)")
print("------------------")
def _augment_small_1():
batch_loader = multicore.BatchLoader(load_images(n_batches=1),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
i = 0
while True:
batch = bg_augmenter.get_batch()
if batch is None:
break
i += 1
for i in range(100):
_augment_small_1()
print("Finished (%d/%d)." % (i + 1, 100))
print("------------------")
print("Many very small runs, separate function (batches=2)")
print("------------------")
def _augment_small_2():
batch_loader = multicore.BatchLoader(load_images(n_batches=2),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
i = 0
while True:
batch = bg_augmenter.get_batch()
if batch is None:
break
i += 1
for i in range(100):
_augment_small_2()
print("Finished (%d/%d)." % (i + 1, 100))
print("------------------")
print(
"Many very small runs, separate function, incomplete fetching (batches=2)"
)
print("------------------")
def _augment_small_3():
batch_loader = multicore.BatchLoader(load_images(n_batches=2),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
batch = bg_augmenter.get_batch()
for i in range(100):
_augment_small_3()
print("Finished (%d/%d)." % (i + 1, 100))
#for augseq_i in [augseq, augseq_slow]:
print("------------------")
print(
"Many very small runs, separate function, incomplete fetching (batches=10)"
)
print("------------------")
def _augment_small_4():
batch_loader = multicore.BatchLoader(load_images(n_batches=10),
queue_size=100)
bg_augmenter = multicore.BackgroundAugmenter(
batch_loader, augseq_i)
batch = bg_augmenter.get_batch()
#bg_augmenter.terminate()
for i in range(100):
_augment_small_4()
print("Finished (%d/%d)." % (i + 1, 100))
def example_background_classes():
print("Example: Background Augmentation via Classes")
import imgaug as ia
import imgaug.multicore as multicore
from imgaug import augmenters as iaa
import numpy as np
from skimage import data
# Example augmentation sequence to run in the background.
augseq = iaa.Sequential([
iaa.Fliplr(0.5),
iaa.CoarseDropout(p=0.1, size_percent=0.1)
])
# A generator that loads batches from the hard drive.
def load_batches():
# Here, load 10 batches of size 4 each.
# You can also load an infinite amount of batches, if you don't train
# in epochs.
batch_size = 4
nb_batches = 10
# Here, for simplicity we just always use the same image.
astronaut = data.astronaut()
astronaut = ia.imresize_single_image(astronaut, (64, 64))
for i in range(nb_batches):
# A list containing all images of the batch.
batch_images = []
# A list containing IDs per image. This is not necessary for the
# background augmentation and here just used to showcase that you
# can transfer additional information.
batch_data = []
# Add some images to the batch.
for b in range(batch_size):
batch_images.append(astronaut)
batch_data.append((i, b))
# Create the batch object to send to the background processes.
batch = ia.Batch(
images=np.array(batch_images, dtype=np.uint8),
data=batch_data
)
yield batch
# background augmentation consists of two components:
# (1) BatchLoader, which runs in a Thread and calls repeatedly a user-defined
# function (here: load_batches) to load batches (optionally with keypoints
# and additional information) and sends them to a queue of batches.
# (2) BackgroundAugmenter, which runs several background processes (on other
# CPU cores). Each process takes batches from the queue defined by (1),
# augments images/keypoints and sends them to another queue.
# The main process can then read augmented batches from the queue defined
# by (2).
batch_loader = multicore.BatchLoader(load_batches)
bg_augmenter = multicore.BackgroundAugmenter(batch_loader, augseq)
# Run until load_batches() returns nothing anymore. This also allows infinite
# training.
while True:
print("Next batch...")
batch = bg_augmenter.get_batch()
if batch is None:
print("Finished epoch.")
break
images_aug = batch.images_aug
print("Image IDs: ", batch.data)
ia.imshow(np.hstack(list(images_aug)))
batch_loader.terminate()
bg_augmenter.terminate()
