Supports Equalize Augmentation. (#543)

* add equalize augmentation in mmcv

* delete unnecessary

* reformat

* remove clip in implementing equalize, and add uint test with case step=0

* remove clip in implementing equalize, and add uint test with case step=0

* add comments for unit test

* rename function name as imequalize

mmcv/image/__init__.py

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

mmcv/image/photometric.py

@@ -123,3 +123,46 @@ def adjust_color(img, alpha=1, beta=None, gamma=0):
         # is needed.
         colored_img = np.clip(colored_img, 0, 255)
     return colored_img
+
+
+def imequalize(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
+        if not step:
+            lut = np.array(range(256))
+        else:
+            # Compute the cumulative sum, shifted by step // 2
+            # and then normalized by step.
+            lut = (np.cumsum(histo) + (step // 2)) // step
+            # Shift lut, prepending with 0.
+            lut = np.concatenate([[0], lut[:-1]], 0)
+        # If step is zero, return the original image.
+        # Otherwise, index from lut.
+        return np.where(np.equal(step, 0), im, lut[im])
+
+    # 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

@@ -106,3 +106,32 @@ class TestPhotometric:
         assert_array_equal(
             np.round(mmcv.adjust_color(img, 0.8, -0.6, gamma=-0.6)),
             np.round(np.clip(img * 0.8 - 0.6 * img_r - 0.6, 0, 255)))
+
+    def test_imequalize(self, nb_rand_test=100):
+
+        def _imequalize(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.imequalize(img)
+        assert_array_equal(equalized_img, _imequalize(img))
+
+        # test equalize with case step=0
+        img = np.array([[0, 0, 0], [120, 120, 120], [255, 255, 255]],
+                       dtype=np.uint8)
+        img = np.stack([img, img, img], axis=-1)
+        assert_array_equal(mmcv.imequalize(img), 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.imequalize(img)
+            assert_array_equal(equalized_img, _imequalize(img))