aleju · aleju · Sep 18, 2019 · Sep 16, 2019 · Sep 16, 2019 · Sep 18, 2019
diff --git a/imgaug/augmenters/pooling.py b/imgaug/augmenters/pooling.py
@@ -25,12 +25,32 @@
 from abc import ABCMeta, abstractmethod
 
 import six
+import numpy as np
 
 from . import meta
 import imgaug as ia
 from .. import parameters as iap
 
 
+def _compute_shape_after_pooling(image_shape, ksize_h, ksize_w):
+    height, width = image_shape[0:2]
+
+    if height == 0:
+        height = 1
+    if width == 0:
+        width = 1
+
+    if height % ksize_h > 0:
+        height += ksize_h - (height % ksize_h)
+    if width % ksize_w > 0:
+        width += ksize_w - (width % ksize_w)
+
+    return tuple([
+        height//ksize_h,
+        width//ksize_w,
+    ] + list(image_shape[2:]))
+
+
 @six.add_metaclass(ABCMeta)
 class _AbstractPoolingBase(meta.Augmenter):
     # TODO add floats as ksize denoting fractions of image sizes
@@ -62,7 +82,10 @@ def _draw_samples(self, augmentables, random_state):
         else:
             kernel_sizes_w = self.kernel_size[1].draw_samples(
                 (nb_images,), random_state=rss[1])
-        return kernel_sizes_h, kernel_sizes_w
+        return (
+            np.clip(kernel_sizes_h, 1, None),
+            np.clip(kernel_sizes_w, 1, None)
+        )
 
     def _augment_images(self, images, random_state, parents, hooks):
         if not self.keep_size:
@@ -75,24 +98,71 @@ def _augment_images(self, images, random_state, parents, hooks):
         for i, (image, ksize_h, ksize_w) in gen:
             if ksize_h >= 2 or ksize_w >= 2:
                 image_pooled = self._pool_image(
-                    image,
-                    max(ksize_h, 1), max(ksize_w, 1)
-                )
+                    image, ksize_h, ksize_w)
                 if self.keep_size:
                     image_pooled = ia.imresize_single_image(
                         image_pooled, image.shape[0:2])
                 images[i] = image_pooled
 
         return images
 
+    def _augment_heatmaps(self, heatmaps, random_state, parents, hooks):
+        return self._augment_hms_and_segmaps(heatmaps, random_state)
+
+    def _augment_segmentation_maps(self, segmaps, random_state, parents, hooks):
+        return self._augment_hms_and_segmaps(segmaps, random_state)
+
+    def _augment_hms_and_segmaps(self, augmentables, random_state):
+        if self.keep_size:
+            return augmentables
+
+        kernel_sizes_h, kernel_sizes_w = self._draw_samples(
+            augmentables, random_state)
+
+        gen = zip(augmentables, kernel_sizes_h, kernel_sizes_w)
+        for augmentable, ksize_h, ksize_w in gen:
+            if ksize_h >= 2 or ksize_w >= 2:
+                # we only update the shape of the underlying image here,
+                # because the library can handle heatmaps/segmaps that are
+                # larger/smaller than the corresponding image
+                new_shape = _compute_shape_after_pooling(
+                    augmentable.shape, ksize_h, ksize_w)
+
+                augmentable.shape = new_shape
+
+        return augmentables
+
+    def _augment_keypoints(self, keypoints_on_images, random_state, parents,
+                           hooks):
+        if self.keep_size:
+            return keypoints_on_images
+
+        kernel_sizes_h, kernel_sizes_w = self._draw_samples(
+            keypoints_on_images, random_state)
+
+        gen = enumerate(zip(keypoints_on_images, kernel_sizes_h,
+                            kernel_sizes_w))
+        for i, (kpsoi, ksize_h, ksize_w) in gen:
+            if ksize_h >= 2 or ksize_w >= 2:
+                new_shape = _compute_shape_after_pooling(
+                    kpsoi.shape, ksize_h, ksize_w)
+
+                keypoints_on_images[i] = kpsoi.on(new_shape)
+
+        return keypoints_on_images
+
+    def _augment_polygons(self, polygons_on_images, random_state, parents,
+                          hooks):
+        return self._augment_polygons_as_keypoints(
+            polygons_on_images, random_state, parents, hooks)
+
     def get_parameters(self):
         return [self.kernel_size, self.keep_size]
 
 
 # TODO rename kernel size parameters in all augmenters to kernel_size
 # TODO add per_channel
 # TODO add upscaling interpolation mode?
-# TODO add dtype support
 class AveragePooling(_AbstractPoolingBase):
     """
     Apply average pooling to images.
@@ -102,15 +172,19 @@ class AveragePooling(_AbstractPoolingBase):
     size. Optionally, the augmenter will automatically re-upscale the image
     to the input size (by default this is activated).
 
-    This augmenter does not affect heatmaps, segmentation maps or
-    coordinates-based augmentables (e.g. keypoints, bounding boxes, ...).
-
     Note that this augmenter is very similar to ``AverageBlur``.
     ``AverageBlur`` applies averaging within windows of given kernel size
     *without* striding, while ``AveragePooling`` applies striding corresponding
     to the kernel size, with optional upscaling afterwards. The upscaling
     is configured to create "pixelated"/"blocky" images by default.
 
+    .. note ::
+
+        During heatmap or segmentation map augmentation, the respective
+        arrays are not changed, only the shapes of the underlying images
+        are updated. This is because imgaug can handle maps/maks that are
+        larger/smaller than their corresponding image.
+
     dtype support::
 
         See :func:`imgaug.imgaug.avg_pool`.
@@ -200,7 +274,7 @@ def __init__(self, kernel_size, keep_size=True,
     def _pool_image(self, image, kernel_size_h, kernel_size_w):
         return ia.avg_pool(
             image,
-            (max(kernel_size_h, 1), max(kernel_size_w, 1))
+            (kernel_size_h, kernel_size_w)
         )
 
 
@@ -213,10 +287,14 @@ class MaxPooling(_AbstractPoolingBase):
     size. Optionally, the augmenter will automatically re-upscale the image
     to the input size (by default this is activated).
 
-    The maximum within each pixel window is always taken channelwise.
+    The maximum within each pixel window is always taken channelwise..
 
-    This augmenter does not affect heatmaps, segmentation maps or
-    coordinates-based augmentables (e.g. keypoints, bounding boxes, ...).
+    .. note ::
+
+        During heatmap or segmentation map augmentation, the respective
+        arrays are not changed, only the shapes of the underlying images
+        are updated. This is because imgaug can handle maps/maks that are
+        larger/smaller than their corresponding image.
 
     dtype support::
 
@@ -309,7 +387,7 @@ def _pool_image(self, image, kernel_size_h, kernel_size_w):
         #      to reflection padding
         return ia.max_pool(
             image,
-            (max(kernel_size_h, 1), max(kernel_size_w, 1))
+            (kernel_size_h, kernel_size_w)
         )
 
 
@@ -324,8 +402,12 @@ class MinPooling(_AbstractPoolingBase):
 
     The minimum within each pixel window is always taken channelwise.
 
-    This augmenter does not affect heatmaps, segmentation maps or
-    coordinates-based augmentables (e.g. keypoints, bounding boxes, ...).
+    .. note ::
+
+        During heatmap or segmentation map augmentation, the respective
+        arrays are not changed, only the shapes of the underlying images
+        are updated. This is because imgaug can handle maps/maks that are
+        larger/smaller than their corresponding image.
 
     dtype support::
 
@@ -418,7 +500,7 @@ def _pool_image(self, image, kernel_size_h, kernel_size_w):
         #      to reflection padding
         return ia.min_pool(
             image,
-            (max(kernel_size_h, 1), max(kernel_size_w, 1))
+            (kernel_size_h, kernel_size_w)
         )
 
 
@@ -433,8 +515,12 @@ class MedianPooling(_AbstractPoolingBase):
 
     The median within each pixel window is always taken channelwise.
 
-    This augmenter does not affect heatmaps, segmentation maps or
-    coordinates-based augmentables (e.g. keypoints, bounding boxes, ...).
+    .. note ::
+
+        During heatmap or segmentation map augmentation, the respective
+        arrays are not changed, only the shapes of the underlying images
+        are updated. This is because imgaug can handle maps/maks that are
+        larger/smaller than their corresponding image.
 
     dtype support::
 
@@ -527,5 +613,5 @@ def _pool_image(self, image, kernel_size_h, kernel_size_w):
         #      to reflection padding
         return ia.median_pool(
             image,
-            (max(kernel_size_h, 1), max(kernel_size_w, 1))
+            (kernel_size_h, kernel_size_w)
         )
diff --git a/imgaug/imgaug.py b/imgaug/imgaug.py
@@ -1678,6 +1678,13 @@ def pad(arr, top=0, right=0, bottom=0, left=0, mode="constant", cval=0):
 
         cval = max(min(cval, max_value), min_value)
 
+        # Note that copyMakeBorder() hangs/runs endlessly if arr has an
+        # axis of size 0 and mode is "reflect".
+        # Numpy also complains in these cases if mode is not "constant".
+        has_zero_sized_axis = any([axis == 0 for axis in arr.shape])
+        if has_zero_sized_axis:
+            mode = "constant"
+
         mapping_mode_np_to_cv2 = {
             "constant": cv2.BORDER_CONSTANT,
             "edge": cv2.BORDER_REPLICATE,
@@ -2103,6 +2110,13 @@ def pool(arr, block_size, func, pad_mode="constant", pad_cval=0,
         pad_cval = cval
 
     _assert_two_or_three_dims(arr)
+    channel_axis_is_zero = (arr.ndim == 3 and arr.shape[-1] == 0)
+
+    # block_reduce() crashes if channel axis is 0. It could probably be
+    # squeezed away, but then how to add a zero-sized axis later on?
+    assert not channel_axis_is_zero, (
+        "Cannot pool a 3d-array with 0 channels. Got shape %s." % (arr.shape,))
+
     is_valid_int = is_single_integer(block_size) and block_size >= 1
     is_valid_tuple = is_iterable(block_size) and len(block_size) in [2, 3] \
         and [is_single_integer(val) and val >= 1 for val in block_size]
@@ -2128,9 +2142,10 @@ def pool(arr, block_size, func, pad_mode="constant", pad_cval=0,
     )
 
     input_dtype = arr.dtype
+
     arr_reduced = skimage.measure.block_reduce(arr, tuple(block_size), func,
                                                cval=cval)
-    if preserve_dtype and arr_reduced.dtype.type != input_dtype:
+    if preserve_dtype and arr_reduced.dtype.name != input_dtype.name:
         arr_reduced = arr_reduced.astype(input_dtype)
     return arr_reduced