Change output buffer size in augment_batches()

CHANGELOG.md

@@ -434,6 +434,8 @@
   the standard that all pixels are given with subpixel accuracy and therefore
   any whole pixel with a maximum should denote the coordinates of that
   pixel's center. #413
+* Changed default `output_buffer_size` in `Augmenter.augment_batches()` from
+  "unlimited" to `10*C`, where `C` is the number of logical CPU cores. #417
 
 
 ## Improved Segmentation Map Augmentation #302

imgaug/augmenters/edges.py

@@ -37,7 +37,7 @@
 # TODO this should be placed in some other file than edges.py as it could be
 #      re-used wherever a binary image is the result
 @six.add_metaclass(ABCMeta)
-class BinaryImageColorizerIf(object):
+class IBinaryImageColorizer(object):
     @abstractmethod
     def colorize(self, image_binary, image_original, nth_image, random_state):
         """
@@ -66,7 +66,7 @@ def colorize(self, image_binary, image_original, nth_image, random_state):
 
 
 # TODO see above, this should be moved to another file
-class RandomColorsBinaryImageColorizer(BinaryImageColorizerIf):
+class RandomColorsBinaryImageColorizer(IBinaryImageColorizer):
     """
     Colorizer using two randomly sampled foreground/background colors.
 

imgaug/augmenters/meta.py

@@ -287,10 +287,19 @@ def __init__(self, name=None, deterministic=False, random_state=None):
     def augment_batches(self, batches, hooks=None, background=False):
         """Augment multiple batches.
 
-        In contrast to other augment functions, this function **yields**
-        batches instead of just returning a full list. This is more suited
-        for most training loops. It also supports augmentation on multiple
-        cpu cores, activated via the `background` flag.
+        In contrast to other ``augment_*`` method, this one **yields**
+        batches instead of returning a full list. This is more suited
+        for most training loops.
+
+        This method also also supports augmentation on multiple cpu cores,
+        activated via the `background` flag. If the `background` flag
+        is activated, an instance of :class:`imgaug.multicore.Pool` will
+        be spawned using all available logical CPU cores and an
+        ``output_buffer_size`` of ``C*10``, where ``C`` is the number of
+        logical CPU cores. I.e. a maximum of ``C*10`` batches will be somewhere
+        in the augmentation pipeline (or waiting to be retrieved by downstream
+        functions) before this method temporarily stops the loading of new
+        batches from `batches`.
 
         Parameters
         ----------
@@ -462,7 +471,12 @@ def load_batches():
                     yield batch_normalized
 
             with multicore.Pool(self) as pool:
-                for batch_aug in pool.imap_batches(load_batches()):
+                # pylint:disable=protected-access
+                # note that pool.processes is None here
+                output_buffer_size = pool.pool._processes * 10
+
+                for batch_aug in pool.imap_batches(
+                        load_batches(), output_buffer_size=output_buffer_size):
                     idx = batch_aug.data[0]
                     assert idx in id_to_batch_orig, (
                         "Got idx %d from Pool, which is not known." % (

imgaug/augmenters/segmentation.py

@@ -580,7 +580,7 @@ def __init__(self, points_sampler, p_replace=1.0, max_size=128,
         super(Voronoi, self).__init__(
             name=name, deterministic=deterministic, random_state=random_state)
 
-        assert isinstance(points_sampler, PointsSamplerIf), (
+        assert isinstance(points_sampler, IPointsSampler), (
             "Expected 'points_sampler' to be an instance of PointsSamplerIf, "
             "got %s." % (type(points_sampler),))
         self.points_sampler = points_sampler
@@ -1083,7 +1083,7 @@ def __init__(self, n_rows_frac, n_cols_frac, p_drop_points=0.4,
 
 
 @six.add_metaclass(ABCMeta)
-class PointsSamplerIf(object):
+class IPointsSampler(object):
     """Interface for all point samplers.
 
     Point samplers return coordinate arrays of shape ``Nx2``.
@@ -1141,7 +1141,7 @@ def _verify_sample_points_images(images):
                 images.ndim, images.shape))
 
 
-class RegularGridPointsSampler(PointsSamplerIf):
+class RegularGridPointsSampler(IPointsSampler):
     """Sampler that generates a regular grid of coordinates on an image.
 
     'Regular grid' here means that on each axis all coordinates have the
@@ -1258,7 +1258,7 @@ def __str__(self):
         return self.__repr__()
 
 
-class RelativeRegularGridPointsSampler(PointsSamplerIf):
+class RelativeRegularGridPointsSampler(IPointsSampler):
     """Regular grid coordinate sampler; places more points on larger images.
 
     This is similar to ``RegularGridPointSampler``, but the number of rows
@@ -1354,7 +1354,7 @@ def __str__(self):
         return self.__repr__()
 
 
-class DropoutPointsSampler(PointsSamplerIf):
+class DropoutPointsSampler(IPointsSampler):
     """Remove a defined fraction of sampled points.
 
     Parameters
@@ -1397,7 +1397,7 @@ class DropoutPointsSampler(PointsSamplerIf):
     """
 
     def __init__(self, other_points_sampler, p_drop):
-        assert isinstance(other_points_sampler, PointsSamplerIf), (
+        assert isinstance(other_points_sampler, IPointsSampler), (
             "Expected to get an instance of PointsSamplerIf as argument "
             "'other_points_sampler', got type %s." % (
                 type(other_points_sampler),))
@@ -1482,7 +1482,7 @@ def __str__(self):
         return self.__repr__()
 
 
-class UniformPointsSampler(PointsSamplerIf):
+class UniformPointsSampler(IPointsSampler):
     """Sample points uniformly on images.
 
     This point sampler generates `n_points` points per image. The x- and
@@ -1563,7 +1563,7 @@ def __str__(self):
         return self.__repr__()
 
 
-class SubsamplingPointsSampler(PointsSamplerIf):
+class SubsamplingPointsSampler(IPointsSampler):
     """Ensure that the number of sampled points is below a maximum.
 
     This point sampler will sample points from another sampler and
@@ -1598,7 +1598,7 @@ class SubsamplingPointsSampler(PointsSamplerIf):
     """
 
     def __init__(self, other_points_sampler, n_points_max):
-        assert isinstance(other_points_sampler, PointsSamplerIf), (
+        assert isinstance(other_points_sampler, IPointsSampler), (
             "Expected to get an instance of PointsSamplerIf as argument "
             "'other_points_sampler', got type %s." % (
                 type(other_points_sampler),))

test/augmenters/test_segmentation.py

@@ -498,7 +498,7 @@ def test_interpolation(self):
         assert mock_imresize.call_args_list[0][1]["interpolation"] == "cubic"
 
     def test_point_sampler_called(self):
-        class LoggedPointSampler(iaa.PointsSamplerIf):
+        class LoggedPointSampler(iaa.IPointsSampler):
             def __init__(self, other):
                 self.other = other
                 self.call_count = 0
@@ -516,7 +516,7 @@ def sample_points(self, images, random_state):
         assert sampler.call_count == 1
 
     def test_point_sampler_returns_no_points_integrationtest(self):
-        class NoPointsPointSampler(iaa.PointsSamplerIf):
+        class NoPointsPointSampler(iaa.IPointsSampler):
             def sample_points(self, images, random_state):
                 return [np.zeros((0, 2), dtype=np.float32)]
 
@@ -1053,7 +1053,7 @@ def test_conversion_to_string(self):
         assert sampler.__str__() == sampler.__repr__() == expected
 
 
-class _FixedPointsSampler(iaa.PointsSamplerIf):
+class _FixedPointsSampler(iaa.IPointsSampler):
     def __init__(self, points):
         self.points = np.float32(np.copy(points))
         self.last_random_state = None