Supports Equalize Augmentation.

mmcv/image/__init__.py

@@ -7,8 +7,8 @@
                         imtranslate, rescale_size)
 from .io import imfrombytes, imread, imwrite, supported_backends, use_backend
 from .misc import tensor2imgs
-from .photometric import (imdenormalize, iminvert, imnormalize, imnormalize_,
-                          posterize, solarize)
+from .photometric import (equalize, imdenormalize, iminvert, imnormalize,
+                          imnormalize_, posterize, solarize)
 
 __all__ = [
     'bgr2gray', 'bgr2hls', 'bgr2hsv', 'bgr2rgb', 'gray2bgr', 'gray2rgb',
@@ -18,5 +18,5 @@
     'imwrite', 'supported_backends', 'use_backend', 'imdenormalize',
     'imnormalize', 'imnormalize_', 'iminvert', 'posterize', 'solarize',
     'rgb2ycbcr', 'bgr2ycbcr', 'ycbcr2rgb', 'ycbcr2bgr', 'tensor2imgs',
-    'imshear', 'imtranslate'
+    'imshear', 'imtranslate', 'equalize'
 ]

mmcv/image/photometric.py

@@ -91,3 +91,51 @@ def posterize(img, bits):
     shift = 8 - bits
     img = np.left_shift(np.right_shift(img, shift), shift)
     return img
+
+
+def equalize(img):
+    """Equalize the image histogram. 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.
+
+    Args:
+        img (ndarray): Image to be equalized.
+
+    Returns:
+        ndarray: The equalized image.
+    """
+
+    def _scale_channel(im, c):
+        """Scale the data in the corresponding channel."""
+        im = im[:, :, c]
+        # Compute the histogram of the image channel.
+        histo = np.histogram(im, 256, (0, 255))[0]
+        # For computing the step, filter out the nonzeros.
+        nonzero_histo = histo[histo > 0]
+        step = (np.sum(nonzero_histo) - nonzero_histo[-1]) // 255
+
+        def _build_lut(histo, step, eps=1e-5):
+            # Compute the cumulative sum, shifted by step // 2
+            # and then normalized by step.
+            if step:
+                lut = (np.cumsum(histo) + (step // 2)) // step
+            else:
+                lut = (np.cumsum(histo) + (step // 2)) // (step + eps)
+            # Shift lut, prepending with 0.
+            lut = np.concatenate([[0], lut[:-1]], 0)
+            # Clip the counts to be in range.
+            return np.clip(lut, 0, 255)
+
+        # If step is zero, return the original image. Otherwise,
+        # build lut from the full histogram and step, and then
+        # index from it.
+        result = np.where(np.equal(step, 0), im, _build_lut(histo, step)[im])
+        return result
+
+    # Scales each channel independently and then stacks
+    # the result.
+    s1 = _scale_channel(img, 0)
+    s2 = _scale_channel(img, 1)
+    s3 = _scale_channel(img, 2)
+    equalized_img = np.stack([s1, s2, s3], axis=-1)
+    return equalized_img

tests/test_image/test_photometric.py

@@ -75,3 +75,26 @@ def test_posterize(self):
         img_r = np.array([[0, 128, 224], [0, 96, 224], [0, 128, 224]],
                          dtype=np.uint8)
         assert_array_equal(mmcv.posterize(img, 3), img_r)
+
+    def test_equalize(self, nb_rand_test=100):
+
+        def _equalize(img):
+            # equalize the image using PIL.ImageOps.equalize
+            from PIL import ImageOps, Image
+            img = Image.fromarray(img)
+            equalized_img = np.asarray(ImageOps.equalize(img))
+            return equalized_img
+
+        img = np.array([[0, 128, 255], [1, 127, 254], [2, 129, 253]],
+                       dtype=np.uint8)
+        img = np.stack([img, img, img], axis=-1)
+        equalized_img = mmcv.equalize(img)
+        assert_array_equal(equalized_img, _equalize(img))
+
+        # test equalize with randomly sampled image.
+        for _ in range(nb_rand_test):
+            img = np.clip(
+                np.random.uniform(0, 1, (1000, 1200, 3)) * 260, 0,
+                255).astype(np.uint8)
+            equalized_img = mmcv.equalize(img)
+            assert_array_equal(equalized_img, _equalize(img))