Merge pull request #480 from aleju/equalize

Add Equalize

changelogs/master/added/20191103_equalize.md

@@ -0,0 +1,6 @@
+# Equalize #480
+
+* Added `imgaug.augmenters.contrast.equalize`, similar to
+  `PIL.ImageOps.equalize`.
+* Added `imgaug.augmenters.contrast.equalize_`.
+* Added `imgaug.augmenters.contrast.Equalize`.

imgaug/augmenters/contrast.py

@@ -37,6 +37,9 @@
 from ..augmentables import batches as iabatches
 
 
+_EQUALIZE_USE_PIL_BELOW = 64*64  # H*W
+
+
 class _ContrastFuncWrapper(meta.Augmenter):
     def __init__(self, func, params1d, per_channel, dtypes_allowed=None,
                  dtypes_disallowed=None,
@@ -431,6 +434,157 @@ def adjust_contrast_linear(arr, alpha):
         return image_aug
 
 
+def equalize(image, mask=None):
+    """Equalize the image histogram.
+
+    See :func:`equalize_` for details.
+
+    This function is identical in inputs and outputs to
+    :func:`PIL.ImageOps.equalize`.
+
+    dtype support::
+
+        See :func:`imgaug.augmenters.contrast.equalize_`.
+
+    Parameters
+    ----------
+    image : ndarray
+        ``uint8`` ``(H,W,[C])`` image to equalize.
+
+    mask : None or ndarray, optional
+        An optional mask. If given, only the pixels selected by the mask are
+        included in the analysis.
+
+    Returns
+    -------
+    ndarray
+        Equalized image.
+
+    """
+    # internally used method works in-place by default and hence needs a copy
+    size = image.size
+    if size == 0:
+        return np.copy(image)
+    elif size >= _EQUALIZE_USE_PIL_BELOW:
+        image = np.copy(image)
+    return equalize_(image, mask)
+
+
+def equalize_(image, mask=None):
+    """Equalize the image histogram in-place.
+
+    This function applies a non-linear mapping to the input image, in order
+    to create a uniform distribution of grayscale values in the output image.
+
+    This function is identical in inputs and outputs to
+    :func:`PIL.ImageOps.equalize`, except that it is allowed to modify the
+    input image in-place.
+
+    dtype support::
+
+        * ``uint8``: yes; fully tested
+        * ``uint16``: no
+        * ``uint32``: no
+        * ``uint64``: no
+        * ``int8``: no
+        * ``int16``: no
+        * ``int32``: no
+        * ``int64``: no
+        * ``float16``: no
+        * ``float32``: no
+        * ``float64``: no
+        * ``float128``: no
+        * ``bool``: no
+
+    Parameters
+    ----------
+    image : ndarray
+        ``uint8`` ``(H,W,[C])`` image to equalize.
+
+    mask : None or ndarray, optional
+        An optional mask. If given, only the pixels selected by the mask are
+        included in the analysis.
+
+    Returns
+    -------
+    ndarray
+        Equalized image. *Might* have been modified in-place.
+
+    """
+    nb_channels = 1 if image.ndim == 2 else image.shape[-1]
+    if nb_channels not in [1, 3]:
+        result = [equalize_(image[:, :, c]) for c in np.arange(nb_channels)]
+        return np.stack(result, axis=-1)
+
+    assert image.dtype.name == "uint8", (
+        "Expected image of dtype uint8, got dtype %s." % (image.dtype.name,))
+    if mask is not None:
+        assert mask.ndim == 2, (
+            "Expected 2-dimensional mask, got shape %s." % (mask.shape,))
+        assert mask.dtype.name == "uint8", (
+            "Expected mask of dtype uint8, got dtype %s." % (mask.dtype.name,))
+
+    size = image.size
+    if size == 0:
+        return image
+    elif nb_channels == 3 and size < _EQUALIZE_USE_PIL_BELOW:
+        return _equalize_pil(image, mask)
+    return _equalize_no_pil_(image, mask)
+
+
+# note that this is supposed to be a non-PIL reimplementation of PIL's
+# equalize, which produces slightly different results from cv2.equalizeHist()
+def _equalize_no_pil_(image, mask=None):
+    flags = image.flags
+    if not flags["OWNDATA"]:
+        image = np.copy(image)
+    if not flags["C_CONTIGUOUS"]:
+        image = np.ascontiguousarray(image)
+
+    nb_channels = 1 if image.ndim == 2 else image.shape[-1]
+    lut = np.empty((1, 256, nb_channels), dtype=np.int32)
+
+    for c_idx in range(nb_channels):
+        if image.ndim == 2:
+            image_c = image[:, :, np.newaxis]
+        else:
+            image_c = image[:, :, c_idx:c_idx+1]
+        histo = cv2.calcHist([image_c], [0], mask, [256], [0, 256])
+        if len(histo.nonzero()[0]) <= 1:
+            lut[0, :, c_idx] = np.arange(256).astype(np.int32)
+            continue
+
+        step = np.sum(histo[:-1]) // 255
+        if not step:
+            lut[0, :, c_idx] = np.arange(256).astype(np.int32)
+            continue
+
+        n = step // 2
+        cs = np.cumsum(histo)
+        lut[0, 0, c_idx] = n
+        lut[0, 1:, c_idx] = n + cs[0:-1]
+        lut[0, :, c_idx] //= int(step)
+    lut = np.clip(lut, None, 255, out=lut).astype(np.uint8)
+    image = cv2.LUT(image, lut, dst=image)
+    if image.ndim == 2 and image.ndim == 3:
+        return image[..., np.newaxis]
+    return image
+
+
+def _equalize_pil(image, mask=None):
+    import PIL.Image
+    import PIL.ImageOps
+
+    if mask is not None:
+        mask = PIL.Image.fromarray(mask).convert("L")
+    return np.asarray(
+        PIL.ImageOps.equalize(
+            PIL.Image.fromarray(image),
+            mask=mask
+        )
+    )
+
+
 class GammaContrast(_ContrastFuncWrapper):
     """
     Adjust image contrast by scaling pixel values to ``255*((v/255)**gamma)``.
@@ -454,7 +608,7 @@ class GammaContrast(_ContrastFuncWrapper):
             * If a ``StochasticParameter``, then a value will be sampled per
               image from that parameter.
 
-    per_channel :  bool or float, optional
+    per_channel : bool or float, optional
         Whether to use the same value for all channels (``False``) or to
         sample a new value for each channel (``True``). If this value is a
         float ``p``, then for ``p`` percent of all images `per_channel` will
@@ -541,7 +695,7 @@ class SigmoidContrast(_ContrastFuncWrapper):
             * If a ``StochasticParameter``, then a value will be sampled per
               image from that parameter.
 
-    per_channel :  bool or float, optional
+    per_channel : bool or float, optional
         Whether to use the same value for all channels (``False``) or to
         sample a new value for each channel (``True``). If this value is a
         float ``p``, then for ``p`` percent of all images `per_channel` will
@@ -624,7 +778,7 @@ class LogContrast(_ContrastFuncWrapper):
             * If a ``StochasticParameter``, then a value will be sampled per
               image from that parameter.
 
-    per_channel :  bool or float, optional
+    per_channel : bool or float, optional
         Whether to use the same value for all channels (``False``) or to
         sample a new value for each channel (``True``). If this value is a
         float ``p``, then for ``p`` percent of all images `per_channel` will
@@ -696,7 +850,7 @@ class LinearContrast(_ContrastFuncWrapper):
             * If a ``StochasticParameter``, then a value will be sampled per
               image from that parameter.
 
-    per_channel :  bool or float, optional
+    per_channel : bool or float, optional
         Whether to use the same value for all channels (``False``) or to
         sample a new value for each channel (``True``). If this value is a
         float ``p``, then for ``p`` percent of all images `per_channel` will
@@ -748,6 +902,50 @@ def __init__(self, alpha=1, per_channel=False,
         )
 
 
+class Equalize(meta.Augmenter):
+    """Equalize the image histogram.
+
+    This augmenter is identical in inputs and outputs to
+    :func:`PIL.ImageOps.equalize`.
+
+    dtype support::
+
+        See :func:`imgaug.augmenters.contrast.equalize_`.
+
+    Parameters
+    ----------
+    name : None or str, optional
+        See :func:`imgaug.augmenters.meta.Augmenter.__init__`.
+
+    deterministic : bool, optional
+        See :func:`imgaug.augmenters.meta.Augmenter.__init__`.
+
+    random_state : None or int or imgaug.random.RNG or numpy.random.Generator or numpy.random.bit_generator.BitGenerator or numpy.random.SeedSequence or numpy.random.RandomState, optional
+        See :func:`imgaug.augmenters.meta.Augmenter.__init__`.
+
+    Examples
+    --------
+    >>> import imgaug.augmenters as iaa
+    >>> aug = iaa.Equalize()
+
+    Equalize the histograms of all input images.
+
+    """
+    def __init__(self, name=None, deterministic=False, random_state=None):
+        super(Equalize, self).__init__(
+            name=name, deterministic=deterministic, random_state=random_state)
+
+    def _augment_batch(self, batch, random_state, parents, hooks):
+        # pylint: disable=no-self-use
+        if batch.images:
+            for image in batch.images:
+                image[...] = equalize_(image)
+        return batch
+
+    def get_parameters(self):
+        return []
+
+
 # TODO maybe offer the other contrast augmenters also wrapped in this, similar
 #      to CLAHE and HistogramEqualization?
 #      this is essentially tested by tests for CLAHE