wehub-resource-sync
931 lines
33 KiB
Python
931 lines
33 KiB
Python
"""Classes and functions dealing with augmentation on multiple CPU cores."""
|
|
from __future__ import print_function, division, absolute_import
|
|
import sys
|
|
import multiprocessing
|
|
import threading
|
|
import traceback
|
|
import time
|
|
import random
|
|
import platform
|
|
|
|
import numpy as np
|
|
import cv2
|
|
|
|
import imgaug.imgaug as ia
|
|
import imgaug.random as iarandom
|
|
from imgaug.augmentables.batches import Batch, UnnormalizedBatch
|
|
|
|
if sys.version_info[0] == 2:
|
|
# pylint: disable=redefined-builtin, import-error
|
|
import cPickle as pickle
|
|
from Queue import Empty as QueueEmpty, Full as QueueFull
|
|
import socket
|
|
BrokenPipeError = socket.error
|
|
elif sys.version_info[0] == 3:
|
|
import pickle
|
|
from queue import Empty as QueueEmpty, Full as QueueFull
|
|
|
|
|
|
_CONTEXT = None
|
|
|
|
|
|
# Added in 0.4.0.
|
|
def _get_context_method():
|
|
vinfo = sys.version_info
|
|
|
|
# get_context() is only supported in 3.5 and later (same for
|
|
# set_start_method)
|
|
get_context_unsupported = (
|
|
vinfo[0] == 2
|
|
or (vinfo[0] == 3 and vinfo[1] <= 3))
|
|
|
|
method = None
|
|
# Fix random hanging code in NixOS by switching to spawn method,
|
|
# see issue #414
|
|
# TODO This is only a workaround and doesn't really fix the underlying
|
|
# issue. The cause of the underlying issue is currently unknown.
|
|
# Its possible that #535 fixes the issue, though earlier tests
|
|
# indicated that the cause was something else.
|
|
# TODO this might break the semaphore used to prevent out of memory
|
|
# errors
|
|
if "NixOS" in platform.version():
|
|
method = "spawn"
|
|
if get_context_unsupported:
|
|
ia.warn("Detected usage of imgaug.multicore in python <=3.4 "
|
|
"and NixOS. This is known to sometimes cause endlessly "
|
|
"hanging programs when also making use of multicore "
|
|
"augmentation (aka background augmentation). Use "
|
|
"python 3.5 or later to prevent this.")
|
|
elif platform.system() == "Darwin" and vinfo[0:2] == (3, 7):
|
|
# On Mac with python 3.7 there seems to be a problem with matplotlib,
|
|
# resulting in the error "libc++abi.dylib: terminating with uncaught
|
|
# exception of type std::runtime_error: Couldn't close file".
|
|
# The error seems to be due to opened files that get closed in
|
|
# child processes and can be prevented by switching to spawn mode.
|
|
# See https://github.com/matplotlib/matplotlib/issues/15410
|
|
# and https://bugs.python.org/issue33725.
|
|
# It is possible that this problem also affects other python versions,
|
|
# but here it only appeared (consistently) in the 3.7 tests and the
|
|
# reports also seem to be focused around 3.7, suggesting explicitly
|
|
# to update to 3.8.2.
|
|
method = "spawn"
|
|
|
|
if get_context_unsupported:
|
|
return False
|
|
return method
|
|
|
|
|
|
# Added in 0.4.0.
|
|
def _set_context(method):
|
|
# method=False indicates that multiprocessing module (i.e. no context)
|
|
# should be used, e.g. because get_context() is not supported
|
|
globals()["_CONTEXT"] = (
|
|
multiprocessing if method is False
|
|
else multiprocessing.get_context(method))
|
|
|
|
|
|
# Added in 0.4.0.
|
|
def _reset_context():
|
|
globals()["_CONTEXT"] = None
|
|
|
|
|
|
# Added in 0.4.0.
|
|
def _autoset_context():
|
|
_set_context(_get_context_method())
|
|
|
|
|
|
# Added in 0.4.0.
|
|
def _get_context():
|
|
if _CONTEXT is None:
|
|
_autoset_context()
|
|
return _CONTEXT
|
|
|
|
|
|
class Pool(object):
|
|
"""
|
|
Wrapper around ``multiprocessing.Pool`` for multicore augmentation.
|
|
|
|
Parameters
|
|
----------
|
|
augseq : imgaug.augmenters.meta.Augmenter
|
|
The augmentation sequence to apply to batches.
|
|
|
|
processes : None or int, optional
|
|
The number of background workers, similar to the same parameter in
|
|
multiprocessing.Pool. If ``None``, the number of the machine's CPU
|
|
cores will be used (this counts hyperthreads as CPU cores). If this is
|
|
set to a negative value ``p``, then ``P - abs(p)`` will be used,
|
|
where ``P`` is the number of CPU cores. E.g. ``-1`` would use all
|
|
cores except one (this is useful to e.g. reserve one core to feed
|
|
batches to the GPU).
|
|
|
|
maxtasksperchild : None or int, optional
|
|
The number of tasks done per worker process before the process is
|
|
killed and restarted, similar to the same parameter in
|
|
multiprocessing.Pool. If ``None``, worker processes will not be
|
|
automatically restarted.
|
|
|
|
seed : None or int, optional
|
|
The seed to use for child processes. If ``None``, a random seed will
|
|
be used.
|
|
|
|
"""
|
|
# This attribute saves the augmentation sequence for background workers so
|
|
# that it does not have to be resend with every batch. The attribute is set
|
|
# once per worker in the worker's initializer. As each worker has its own
|
|
# process, it is a different variable per worker (though usually should be
|
|
# of equal content).
|
|
_WORKER_AUGSEQ = None
|
|
|
|
# This attribute saves the initial seed for background workers so that for
|
|
# any future batch the batch's specific seed can be derived, roughly via
|
|
# SEED_START+SEED_BATCH. As each worker has its own process, this seed can
|
|
# be unique per worker even though all seemingly use the same constant
|
|
# attribute.
|
|
_WORKER_SEED_START = None
|
|
|
|
def __init__(self, augseq, processes=None, maxtasksperchild=None,
|
|
seed=None):
|
|
# make sure that don't call pool again in a child process
|
|
assert Pool._WORKER_AUGSEQ is None, (
|
|
"_WORKER_AUGSEQ was already set when calling Pool.__init__(). "
|
|
"Did you try to instantiate a Pool within a Pool?")
|
|
assert processes is None or processes != 0, (
|
|
"Expected `processes` to be `None` (\"use as many cores as "
|
|
"available\") or a negative integer (\"use as many as available "
|
|
"MINUS this number\") or an integer>1 (\"use exactly that many "
|
|
"processes\"). Got type %s, value %s instead." % (
|
|
type(processes), str(processes))
|
|
)
|
|
|
|
self.augseq = augseq
|
|
self.processes = processes
|
|
self.maxtasksperchild = maxtasksperchild
|
|
|
|
if seed is not None:
|
|
assert iarandom.SEED_MIN_VALUE <= seed <= iarandom.SEED_MAX_VALUE, (
|
|
"Expected `seed` to be either `None` or a value between "
|
|
"%d and %d. Got type %s, value %s instead." % (
|
|
iarandom.SEED_MIN_VALUE,
|
|
iarandom.SEED_MAX_VALUE,
|
|
type(seed),
|
|
str(seed)
|
|
)
|
|
)
|
|
self.seed = seed
|
|
|
|
# multiprocessing.Pool instance
|
|
self._pool = None
|
|
|
|
# Running counter of the number of augmented batches. This will be
|
|
# used to send indexes for each batch to the workers so that they can
|
|
# augment using SEED_BASE+SEED_BATCH and ensure consistency of applied
|
|
# augmentation order between script runs.
|
|
self._batch_idx = 0
|
|
|
|
@property
|
|
def pool(self):
|
|
"""Return or create the ``multiprocessing.Pool`` instance.
|
|
|
|
This creates a new instance upon the first call and afterwards
|
|
returns that instance (until the property ``_pool`` is set to
|
|
``None`` again).
|
|
|
|
Returns
|
|
-------
|
|
multiprocessing.Pool
|
|
The ``multiprocessing.Pool`` used internally by this
|
|
``imgaug.multicore.Pool``.
|
|
|
|
"""
|
|
if self._pool is None:
|
|
processes = self.processes
|
|
if processes is not None and processes < 0:
|
|
# cpu count returns the number of logical cpu cores, i.e.
|
|
# including hyperthreads could also use
|
|
# os.sched_getaffinity(0) here, which seems to not exist on
|
|
# BSD though.
|
|
# In python 3.4+, there is also os.cpu_count(), which
|
|
# multiprocessing.cpu_count() then redirects to.
|
|
# At least one guy on stackoverflow.com/questions/1006289
|
|
# reported that only os.* existed, not the multiprocessing
|
|
# method.
|
|
# TODO make this also check if os.cpu_count exists as a
|
|
# fallback
|
|
try:
|
|
processes = _get_context().cpu_count() - abs(processes)
|
|
processes = max(processes, 1)
|
|
except (ImportError, NotImplementedError):
|
|
ia.warn(
|
|
"Could not find method multiprocessing.cpu_count(). "
|
|
"This will likely lead to more CPU cores being used "
|
|
"for the background augmentation than originally "
|
|
"intended.")
|
|
processes = None
|
|
|
|
self._pool = _get_context().Pool(
|
|
processes,
|
|
initializer=_Pool_initialize_worker,
|
|
initargs=(self.augseq, self.seed),
|
|
maxtasksperchild=self.maxtasksperchild)
|
|
return self._pool
|
|
|
|
def map_batches(self, batches, chunksize=None):
|
|
"""
|
|
Augment a list of batches.
|
|
|
|
Parameters
|
|
----------
|
|
batches : list of imgaug.augmentables.batches.Batch
|
|
The batches to augment.
|
|
|
|
chunksize : None or int, optional
|
|
Rough indicator of how many tasks should be sent to each worker.
|
|
Increasing this number can improve performance.
|
|
|
|
Returns
|
|
-------
|
|
list of imgaug.augmentables.batches.Batch
|
|
Augmented batches.
|
|
|
|
"""
|
|
self._assert_batches_is_list(batches)
|
|
return self.pool.map(
|
|
_Pool_starworker,
|
|
self._handle_batch_ids(batches),
|
|
chunksize=chunksize)
|
|
|
|
def map_batches_async(self, batches, chunksize=None, callback=None,
|
|
error_callback=None):
|
|
"""
|
|
Augment batches asynchonously.
|
|
|
|
Parameters
|
|
----------
|
|
batches : list of imgaug.augmentables.batches.Batch
|
|
The batches to augment.
|
|
|
|
chunksize : None or int, optional
|
|
Rough indicator of how many tasks should be sent to each worker.
|
|
Increasing this number can improve performance.
|
|
|
|
callback : None or callable, optional
|
|
Function to call upon finish. See ``multiprocessing.Pool``.
|
|
|
|
error_callback : None or callable, optional
|
|
Function to call upon errors. See ``multiprocessing.Pool``.
|
|
|
|
Returns
|
|
-------
|
|
multiprocessing.MapResult
|
|
Asynchonous result. See ``multiprocessing.Pool``.
|
|
|
|
"""
|
|
self._assert_batches_is_list(batches)
|
|
return self.pool.map_async(
|
|
_Pool_starworker,
|
|
self._handle_batch_ids(batches),
|
|
chunksize=chunksize,
|
|
callback=callback,
|
|
error_callback=error_callback)
|
|
|
|
@classmethod
|
|
def _assert_batches_is_list(cls, batches):
|
|
assert isinstance(batches, list), (
|
|
"Expected `batches` to be a list, got type %s. Call "
|
|
"imap_batches() if you use generators.") % (type(batches),)
|
|
|
|
def imap_batches(self, batches, chunksize=1, output_buffer_size=None):
|
|
"""
|
|
Augment batches from a generator.
|
|
|
|
Pattern for output buffer constraint is from
|
|
https://stackoverflow.com/a/47058399.
|
|
|
|
Parameters
|
|
----------
|
|
batches : generator of imgaug.augmentables.batches.Batch
|
|
The batches to augment, provided as a generator. Each call to the
|
|
generator should yield exactly one batch.
|
|
|
|
chunksize : None or int, optional
|
|
Rough indicator of how many tasks should be sent to each worker.
|
|
Increasing this number can improve performance.
|
|
|
|
output_buffer_size : None or int, optional
|
|
Max number of batches to handle *at the same time* in the *whole*
|
|
pipeline (including already augmented batches that are waiting to
|
|
be requested). If the buffer size is reached, no new batches will
|
|
be loaded from `batches` until a produced (i.e. augmented) batch is
|
|
consumed (i.e. requested from this method).
|
|
The buffer is unlimited if this is set to ``None``. For large
|
|
datasets, this should be set to an integer value to avoid filling
|
|
the whole RAM if loading+augmentation happens faster than training.
|
|
|
|
*New in version 0.3.0.*
|
|
|
|
Yields
|
|
------
|
|
imgaug.augmentables.batches.Batch
|
|
Augmented batch.
|
|
|
|
"""
|
|
self._assert_batches_is_generator(batches)
|
|
|
|
# buffer is either None or a Semaphore
|
|
output_buffer_left = _create_output_buffer_left(output_buffer_size)
|
|
|
|
# TODO change this to 'yield from' once switched to 3.3+
|
|
gen = self.pool.imap(
|
|
_Pool_starworker,
|
|
self._ibuffer_batch_loading(
|
|
self._handle_batch_ids_gen(batches),
|
|
output_buffer_left
|
|
),
|
|
chunksize=chunksize)
|
|
|
|
for batch in gen:
|
|
yield batch
|
|
if output_buffer_left is not None:
|
|
output_buffer_left.release()
|
|
|
|
def imap_batches_unordered(self, batches, chunksize=1,
|
|
output_buffer_size=None):
|
|
"""Augment batches from a generator (without preservation of order).
|
|
|
|
Pattern for output buffer constraint is from
|
|
https://stackoverflow.com/a/47058399.
|
|
|
|
Parameters
|
|
----------
|
|
batches : generator of imgaug.augmentables.batches.Batch
|
|
The batches to augment, provided as a generator. Each call to the
|
|
generator should yield exactly one batch.
|
|
|
|
chunksize : None or int, optional
|
|
Rough indicator of how many tasks should be sent to each worker.
|
|
Increasing this number can improve performance.
|
|
|
|
output_buffer_size : None or int, optional
|
|
Max number of batches to handle *at the same time* in the *whole*
|
|
pipeline (including already augmented batches that are waiting to
|
|
be requested). If the buffer size is reached, no new batches will
|
|
be loaded from `batches` until a produced (i.e. augmented) batch is
|
|
consumed (i.e. requested from this method).
|
|
The buffer is unlimited if this is set to ``None``. For large
|
|
datasets, this should be set to an integer value to avoid filling
|
|
the whole RAM if loading+augmentation happens faster than training.
|
|
|
|
*New in version 0.3.0.*
|
|
|
|
Yields
|
|
------
|
|
imgaug.augmentables.batches.Batch
|
|
Augmented batch.
|
|
|
|
"""
|
|
self._assert_batches_is_generator(batches)
|
|
|
|
# buffer is either None or a Semaphore
|
|
output_buffer_left = _create_output_buffer_left(output_buffer_size)
|
|
|
|
gen = self.pool.imap_unordered(
|
|
_Pool_starworker,
|
|
self._ibuffer_batch_loading(
|
|
self._handle_batch_ids_gen(batches),
|
|
output_buffer_left
|
|
),
|
|
chunksize=chunksize
|
|
)
|
|
|
|
for batch in gen:
|
|
yield batch
|
|
if output_buffer_left is not None:
|
|
output_buffer_left.release()
|
|
|
|
@classmethod
|
|
def _assert_batches_is_generator(cls, batches):
|
|
assert ia.is_generator(batches), (
|
|
"Expected `batches` to be generator, got type %s. Call "
|
|
"map_batches() if you use lists.") % (type(batches),)
|
|
|
|
def __enter__(self):
|
|
assert self._pool is None, (
|
|
"Tried to __enter__ a pool that has already been initialized.")
|
|
_ = self.pool # initialize internal multiprocessing pool instance
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.close()
|
|
|
|
def close(self):
|
|
"""Close the pool gracefully."""
|
|
if self._pool is not None:
|
|
self._pool.close()
|
|
self._pool.join()
|
|
self._pool = None
|
|
|
|
def terminate(self):
|
|
"""Terminate the pool immediately."""
|
|
if self._pool is not None:
|
|
self._pool.terminate()
|
|
self._pool.join()
|
|
self._pool = None
|
|
|
|
# TODO why does this function exist if it may only be called after
|
|
# close/terminate and both of these two already call join() themselves
|
|
def join(self):
|
|
"""
|
|
Wait for the workers to exit.
|
|
|
|
This may only be called after first calling
|
|
:func:`~imgaug.multicore.Pool.close` or
|
|
:func:`~imgaug.multicore.Pool.terminate`.
|
|
|
|
"""
|
|
if self._pool is not None:
|
|
self._pool.join()
|
|
|
|
def _handle_batch_ids(self, batches):
|
|
ids = np.arange(self._batch_idx, self._batch_idx + len(batches))
|
|
inputs = list(zip(ids, batches))
|
|
self._batch_idx += len(batches)
|
|
return inputs
|
|
|
|
def _handle_batch_ids_gen(self, batches):
|
|
for batch in batches:
|
|
batch_idx = self._batch_idx
|
|
yield batch_idx, batch
|
|
self._batch_idx += 1
|
|
|
|
@classmethod
|
|
def _ibuffer_batch_loading(cls, batches, output_buffer_left):
|
|
for batch in batches:
|
|
if output_buffer_left is not None:
|
|
output_buffer_left.acquire()
|
|
yield batch
|
|
|
|
|
|
def _create_output_buffer_left(output_buffer_size):
|
|
output_buffer_left = None
|
|
if output_buffer_size:
|
|
assert output_buffer_size > 0, (
|
|
"Expected buffer size to be greater than zero, but got size %d "
|
|
"instead." % (output_buffer_size,))
|
|
output_buffer_left = _get_context().Semaphore(output_buffer_size)
|
|
return output_buffer_left
|
|
|
|
|
|
# This could be a classmethod or staticmethod of Pool in 3.x, but in 2.7 that
|
|
# leads to pickle errors.
|
|
def _Pool_initialize_worker(augseq, seed_start):
|
|
# pylint: disable=invalid-name, protected-access
|
|
|
|
# Not using this seems to have caused infinite hanging in the case
|
|
# of gaussian blur on at least MacOSX.
|
|
# It is also in most cases probably not sensible to use multiple
|
|
# threads while already running augmentation in multiple processes.
|
|
cv2.setNumThreads(0)
|
|
|
|
if seed_start is None:
|
|
# pylint falsely thinks in older versions that
|
|
# multiprocessing.current_process() was not callable, see
|
|
# https://github.com/PyCQA/pylint/issues/1699
|
|
# pylint: disable=not-callable
|
|
process_name = _get_context().current_process().name
|
|
# pylint: enable=not-callable
|
|
|
|
# time_ns() exists only in 3.7+
|
|
if sys.version_info[0] == 3 and sys.version_info[1] >= 7:
|
|
seed_offset = time.time_ns()
|
|
else:
|
|
seed_offset = int(time.time() * 10**6) % 10**6
|
|
seed = hash(process_name) + seed_offset
|
|
_reseed_global_local(seed, augseq)
|
|
Pool._WORKER_SEED_START = seed_start
|
|
Pool._WORKER_AUGSEQ = augseq
|
|
# not sure if really necessary, but shouldn't hurt either
|
|
Pool._WORKER_AUGSEQ.localize_random_state_()
|
|
|
|
|
|
# This could be a classmethod or staticmethod of Pool in 3.x, but in 2.7 that
|
|
# leads to pickle errors.
|
|
def _Pool_worker(batch_idx, batch):
|
|
# pylint: disable=invalid-name, protected-access
|
|
assert ia.is_single_integer(batch_idx), (
|
|
"Expected `batch_idx` to be an integer. Got type %s instead." % (
|
|
type(batch_idx)
|
|
))
|
|
assert isinstance(batch, (UnnormalizedBatch, Batch)), (
|
|
"Expected `batch` to be either an instance of "
|
|
"`imgaug.augmentables.batches.UnnormalizedBatch` or "
|
|
"`imgaug.augmentables.batches.Batch`. Got type %s instead." % (
|
|
type(batch)
|
|
))
|
|
assert Pool._WORKER_AUGSEQ is not None, (
|
|
"Expected `Pool._WORKER_AUGSEQ` to NOT be `None`. Did you manually "
|
|
"call _Pool_worker()?")
|
|
|
|
augseq = Pool._WORKER_AUGSEQ
|
|
# TODO why is this if here? _WORKER_SEED_START should always be set?
|
|
if Pool._WORKER_SEED_START is not None:
|
|
seed = Pool._WORKER_SEED_START + batch_idx
|
|
_reseed_global_local(seed, augseq)
|
|
result = augseq.augment_batch_(batch)
|
|
return result
|
|
|
|
|
|
# could be a classmethod or staticmethod of Pool in 3.x, but in 2.7 that leads
|
|
# to pickle errors starworker is here necessary, because starmap does not exist
|
|
# in 2.7
|
|
def _Pool_starworker(inputs):
|
|
# pylint: disable=invalid-name
|
|
return _Pool_worker(*inputs)
|
|
|
|
|
|
def _reseed_global_local(base_seed, augseq):
|
|
seed_global = _derive_seed(base_seed, -10**9)
|
|
seed_local = _derive_seed(base_seed)
|
|
iarandom.seed(seed_global)
|
|
augseq.seed_(seed_local)
|
|
|
|
|
|
def _derive_seed(base_seed, offset=0):
|
|
return (
|
|
iarandom.SEED_MIN_VALUE
|
|
+ (base_seed + offset)
|
|
% (iarandom.SEED_MAX_VALUE - iarandom.SEED_MIN_VALUE)
|
|
)
|
|
|
|
|
|
class BatchLoader(object):
|
|
"""**Deprecated**. Load batches in the background.
|
|
|
|
Deprecated. Use ``imgaug.multicore.Pool`` instead.
|
|
|
|
Loaded batches can be accesses using :attr:`imgaug.BatchLoader.queue`.
|
|
|
|
Parameters
|
|
----------
|
|
load_batch_func : callable or generator
|
|
Generator or generator function (i.e. function that yields Batch
|
|
objects) or a function that returns a list of Batch objects.
|
|
Background loading automatically stops when the last batch was yielded
|
|
or the last batch in the list was reached.
|
|
|
|
queue_size : int, optional
|
|
Maximum number of batches to store in the queue. May be set higher
|
|
for small images and/or small batches.
|
|
|
|
nb_workers : int, optional
|
|
Number of workers to run in the background.
|
|
|
|
threaded : bool, optional
|
|
Whether to run the background processes using threads (True) or full
|
|
processes (False).
|
|
|
|
"""
|
|
|
|
@ia.deprecated(alt_func="imgaug.multicore.Pool")
|
|
def __init__(self, load_batch_func, queue_size=50, nb_workers=1,
|
|
threaded=True):
|
|
assert queue_size >= 2, (
|
|
"Queue size for BatchLoader must be at least 2, "
|
|
"got %d." % (queue_size,))
|
|
assert nb_workers >= 1, (
|
|
"Number of workers for BatchLoader must be at least 1, "
|
|
"got %d" % (nb_workers,))
|
|
self._queue_internal = multiprocessing.Queue(queue_size//2)
|
|
self.queue = multiprocessing.Queue(queue_size//2)
|
|
self.join_signal = multiprocessing.Event()
|
|
self.workers = []
|
|
self.threaded = threaded
|
|
seeds = iarandom.get_global_rng().generate_seeds_(nb_workers)
|
|
for i in range(nb_workers):
|
|
if threaded:
|
|
worker = threading.Thread(
|
|
target=self._load_batches,
|
|
args=(load_batch_func, self._queue_internal,
|
|
self.join_signal, None)
|
|
)
|
|
else:
|
|
worker = multiprocessing.Process(
|
|
target=self._load_batches,
|
|
args=(load_batch_func, self._queue_internal,
|
|
self.join_signal, seeds[i])
|
|
)
|
|
worker.daemon = True
|
|
worker.start()
|
|
self.workers.append(worker)
|
|
|
|
self.main_worker_thread = threading.Thread(
|
|
target=self._main_worker,
|
|
args=()
|
|
)
|
|
self.main_worker_thread.daemon = True
|
|
self.main_worker_thread.start()
|
|
|
|
def count_workers_alive(self):
|
|
return sum([int(worker.is_alive()) for worker in self.workers])
|
|
|
|
def all_finished(self):
|
|
"""
|
|
Determine whether the workers have finished the loading process.
|
|
|
|
Returns
|
|
-------
|
|
out : bool
|
|
True if all workers have finished. Else False.
|
|
|
|
"""
|
|
return self.count_workers_alive() == 0
|
|
|
|
def _main_worker(self):
|
|
workers_running = self.count_workers_alive()
|
|
|
|
while workers_running > 0 and not self.join_signal.is_set():
|
|
# wait for a new batch in the source queue and load it
|
|
try:
|
|
batch_str = self._queue_internal.get(timeout=0.1)
|
|
if batch_str == "":
|
|
workers_running -= 1
|
|
else:
|
|
self.queue.put(batch_str)
|
|
except QueueEmpty:
|
|
time.sleep(0.01)
|
|
except (EOFError, BrokenPipeError):
|
|
break
|
|
|
|
workers_running = self.count_workers_alive()
|
|
|
|
# All workers have finished, move the remaining entries from internal
|
|
# to external queue
|
|
while True:
|
|
try:
|
|
batch_str = self._queue_internal.get(timeout=0.005)
|
|
if batch_str != "":
|
|
self.queue.put(batch_str)
|
|
except QueueEmpty:
|
|
break
|
|
except (EOFError, BrokenPipeError):
|
|
break
|
|
|
|
self.queue.put(pickle.dumps(None, protocol=-1))
|
|
time.sleep(0.01)
|
|
|
|
@classmethod
|
|
def _load_batches(cls, load_batch_func, queue_internal, join_signal,
|
|
seedval):
|
|
# pylint: disable=broad-except
|
|
if seedval is not None:
|
|
random.seed(seedval)
|
|
np.random.seed(seedval)
|
|
iarandom.seed(seedval)
|
|
|
|
try:
|
|
gen = (
|
|
load_batch_func()
|
|
if not ia.is_generator(load_batch_func)
|
|
else load_batch_func
|
|
)
|
|
for batch in gen:
|
|
assert isinstance(batch, Batch), (
|
|
"Expected batch returned by load_batch_func to "
|
|
"be of class imgaug.Batch, got %s." % (
|
|
type(batch),))
|
|
batch_pickled = pickle.dumps(batch, protocol=-1)
|
|
while not join_signal.is_set():
|
|
try:
|
|
queue_internal.put(batch_pickled, timeout=0.005)
|
|
break
|
|
except QueueFull:
|
|
pass
|
|
if join_signal.is_set():
|
|
break
|
|
except Exception:
|
|
traceback.print_exc()
|
|
finally:
|
|
queue_internal.put("")
|
|
time.sleep(0.01)
|
|
|
|
def terminate(self):
|
|
"""Stop all workers."""
|
|
# pylint: disable=protected-access
|
|
if not self.join_signal.is_set():
|
|
self.join_signal.set()
|
|
# give minimal time to put generated batches in queue and gracefully
|
|
# shut down
|
|
time.sleep(0.01)
|
|
|
|
if self.main_worker_thread.is_alive():
|
|
self.main_worker_thread.join()
|
|
|
|
if self.threaded:
|
|
for worker in self.workers:
|
|
if worker.is_alive():
|
|
worker.join()
|
|
else:
|
|
for worker in self.workers:
|
|
if worker.is_alive():
|
|
worker.terminate()
|
|
worker.join()
|
|
|
|
# wait until all workers are fully terminated
|
|
while not self.all_finished():
|
|
time.sleep(0.001)
|
|
|
|
# empty queue until at least one element can be added and place None
|
|
# as signal that BL finished
|
|
if self.queue.full():
|
|
self.queue.get()
|
|
self.queue.put(pickle.dumps(None, protocol=-1))
|
|
time.sleep(0.01)
|
|
|
|
# clean the queue, this reportedly prevents hanging threads
|
|
while True:
|
|
try:
|
|
self._queue_internal.get(timeout=0.005)
|
|
except QueueEmpty:
|
|
break
|
|
|
|
if not self._queue_internal._closed:
|
|
self._queue_internal.close()
|
|
if not self.queue._closed:
|
|
self.queue.close()
|
|
self._queue_internal.join_thread()
|
|
self.queue.join_thread()
|
|
time.sleep(0.025)
|
|
|
|
def __del__(self):
|
|
if not self.join_signal.is_set():
|
|
self.join_signal.set()
|
|
|
|
|
|
class BackgroundAugmenter(object):
|
|
"""
|
|
**Deprecated**. Augment batches in the background processes.
|
|
|
|
Deprecated. Use ``imgaug.multicore.Pool`` instead.
|
|
|
|
This is a wrapper around the multiprocessing module.
|
|
|
|
Parameters
|
|
----------
|
|
batch_loader : BatchLoader or multiprocessing.Queue
|
|
BatchLoader object that loads the data fed into the
|
|
BackgroundAugmenter, or alternatively a Queue. If a Queue, then it
|
|
must be made sure that a final ``None`` in the Queue signals that the
|
|
loading is finished and no more batches will follow. Otherwise the
|
|
BackgroundAugmenter will wait forever for the next batch.
|
|
|
|
augseq : Augmenter
|
|
An augmenter to apply to all loaded images.
|
|
This may be e.g. a Sequential to apply multiple augmenters.
|
|
|
|
queue_size : int
|
|
Size of the queue that is used to temporarily save the augmentation
|
|
results. Larger values offer the background processes more room
|
|
to save results when the main process doesn't load much, i.e. they
|
|
can lead to smoother and faster training. For large images, high
|
|
values can block a lot of RAM though.
|
|
|
|
nb_workers : 'auto' or int
|
|
Number of background workers to spawn.
|
|
If ``auto``, it will be set to ``C-1``, where ``C`` is the number of
|
|
CPU cores.
|
|
|
|
"""
|
|
|
|
@ia.deprecated(alt_func="imgaug.multicore.Pool")
|
|
def __init__(self, batch_loader, augseq, queue_size=50, nb_workers="auto"):
|
|
assert queue_size > 0, (
|
|
"Expected 'queue_size' to be at least 1, got %d." % (queue_size,))
|
|
self.augseq = augseq
|
|
self.queue_source = (
|
|
batch_loader
|
|
if isinstance(batch_loader, multiprocessing.queues.Queue)
|
|
else batch_loader.queue
|
|
)
|
|
self.queue_result = multiprocessing.Queue(queue_size)
|
|
|
|
if nb_workers == "auto":
|
|
try:
|
|
nb_workers = multiprocessing.cpu_count()
|
|
except (ImportError, NotImplementedError):
|
|
nb_workers = 1
|
|
# try to reserve at least one core for the main process
|
|
nb_workers = max(1, nb_workers - 1)
|
|
else:
|
|
assert nb_workers >= 1, (
|
|
"Expected 'nb_workers' to be \"auto\" or at least 1, "
|
|
"got %d instead." % (nb_workers,))
|
|
|
|
self.nb_workers = nb_workers
|
|
self.workers = []
|
|
self.nb_workers_finished = 0
|
|
|
|
seeds = iarandom.get_global_rng().generate_seeds_(nb_workers)
|
|
for i in range(nb_workers):
|
|
worker = multiprocessing.Process(
|
|
target=self._augment_images_worker,
|
|
args=(augseq, self.queue_source, self.queue_result, seeds[i])
|
|
)
|
|
worker.daemon = True
|
|
worker.start()
|
|
self.workers.append(worker)
|
|
|
|
def all_finished(self):
|
|
return self.nb_workers_finished == self.nb_workers
|
|
|
|
def get_batch(self):
|
|
"""
|
|
Returns a batch from the queue of augmented batches.
|
|
|
|
If workers are still running and there are no batches in the queue,
|
|
it will automatically wait for the next batch.
|
|
|
|
Returns
|
|
-------
|
|
out : None or imgaug.Batch
|
|
One batch or None if all workers have finished.
|
|
|
|
"""
|
|
if self.all_finished():
|
|
return None
|
|
|
|
batch_str = self.queue_result.get()
|
|
batch = pickle.loads(batch_str)
|
|
if batch is not None:
|
|
return batch
|
|
|
|
self.nb_workers_finished += 1
|
|
if self.nb_workers_finished >= self.nb_workers:
|
|
try:
|
|
# remove `None` from the source queue
|
|
self.queue_source.get(timeout=0.001)
|
|
except QueueEmpty:
|
|
pass
|
|
return None
|
|
return self.get_batch()
|
|
|
|
@classmethod
|
|
def _augment_images_worker(cls, augseq, queue_source, queue_result,
|
|
seedval):
|
|
"""
|
|
Augment endlessly images in the source queue.
|
|
|
|
This is a worker function for that endlessly queries the source queue
|
|
(input batches), augments batches in it and sends the result to the
|
|
output queue.
|
|
|
|
"""
|
|
np.random.seed(seedval)
|
|
random.seed(seedval)
|
|
augseq.seed_(seedval)
|
|
iarandom.seed(seedval)
|
|
|
|
loader_finished = False
|
|
|
|
while not loader_finished:
|
|
# wait for a new batch in the source queue and load it
|
|
try:
|
|
batch_str = queue_source.get(timeout=0.1)
|
|
batch = pickle.loads(batch_str)
|
|
if batch is None:
|
|
loader_finished = True
|
|
# put it back in so that other workers know that the
|
|
# loading queue is finished
|
|
queue_source.put(pickle.dumps(None, protocol=-1))
|
|
else:
|
|
batch_aug = augseq.augment_batch_(batch)
|
|
|
|
# send augmented batch to output queue
|
|
batch_str = pickle.dumps(batch_aug, protocol=-1)
|
|
queue_result.put(batch_str)
|
|
except QueueEmpty:
|
|
time.sleep(0.01)
|
|
|
|
queue_result.put(pickle.dumps(None, protocol=-1))
|
|
time.sleep(0.01)
|
|
|
|
def terminate(self):
|
|
"""
|
|
Terminates all background processes immediately.
|
|
|
|
This will also free their RAM.
|
|
|
|
"""
|
|
# pylint: disable=protected-access
|
|
for worker in self.workers:
|
|
if worker.is_alive():
|
|
worker.terminate()
|
|
self.nb_workers_finished = len(self.workers)
|
|
|
|
if not self.queue_result._closed:
|
|
self.queue_result.close()
|
|
time.sleep(0.01)
|
|
|
|
def __del__(self):
|
|
time.sleep(0.1)
|
|
self.terminate()
|