[Feature] Support auto contrast

mmcv/image/__init__.py

@@ -8,9 +8,9 @@
 from .io import imfrombytes, imread, imwrite, supported_backends, use_backend
 from .misc import tensor2imgs
 from .photometric import (adjust_brightness, adjust_color, adjust_contrast,
-                          adjust_sharpness, clahe, imdenormalize, imequalize,
-                          iminvert, imnormalize, imnormalize_, lut_transform,
-                          posterize, solarize)
+                          adjust_sharpness, auto_contrast, clahe,
+                          imdenormalize, imequalize, iminvert, imnormalize,
+                          imnormalize_, lut_transform, posterize, solarize)
 
 __all__ = [
     'bgr2gray', 'bgr2hls', 'bgr2hsv', 'bgr2rgb', 'gray2bgr', 'gray2rgb',
@@ -22,5 +22,5 @@
     'rgb2ycbcr', 'bgr2ycbcr', 'ycbcr2rgb', 'ycbcr2bgr', 'tensor2imgs',
     'imshear', 'imtranslate', 'adjust_color', 'imequalize',
     'adjust_brightness', 'adjust_contrast', 'lut_transform', 'clahe',
-    'adjust_sharpness'
+    'adjust_sharpness', 'auto_contrast'
 ]

mmcv/image/photometric.py

@@ -232,6 +232,62 @@ def adjust_contrast(img, factor=1.):
     return contrasted_img.astype(img.dtype)
 
 
+def auto_contrast(img, cut_off=0):
+    """Auto adjust image contrast.
+
+    This function maximize (normalize) image contrast by first removing cutoff
+    percent of the lightest and darkest pixels from the histogram and remapping
+     the image so that the darkest pixel becomes black (0), and the lightest
+    becomes white (255).
+
+    Args:
+        img (ndarray): Image to be contrasted. BGR order.
+        cut_off (int | float | tuple): The cutoff percent of the lightest and
+            darkest pixels to be removed. If given as tuple, it shall be
+            (low, high). Otherwise, the single value will be used for both.
+            Defaults to 0.
+
+    Returns:
+        ndarray: The contrasted image.
+    """
+
+    def _auto_contrast_channel(im, c, cut_off):
+        im = im[:, :, c]
+        # Compute the histogram of the image channel.
+        histo = np.histogram(im, 256, (0, 255))[0]
+        # Remove cut-off percent pixels from histo
+        histo_sum = np.cumsum(histo)
+        cut_low = histo_sum[-1] * cut_off[0] // 100
+        cut_high = histo_sum[-1] - histo_sum[-1] * cut_off[1] // 100
+        histo_sum = np.clip(histo_sum, cut_low, cut_high) - cut_low
+        histo = np.concatenate([[histo_sum[0]], np.diff(histo_sum)], 0)
+
+        # Compute mapping
+        low, high = np.nonzero(histo)[0][0], np.nonzero(histo)[0][-1]
+        # If all the values have been cut off, return the origin img
+        if low >= high:
+            return im
+        scale = 255.0 / (high - low)
+        offset = -low * scale
+        lut = np.array(range(256))
+        lut = lut * scale + offset
+        lut = np.clip(lut, 0, 255)
+        return lut[im]
+
+    if isinstance(cut_off, (int, float)):
+        cut_off = (cut_off, cut_off)
+    else:
+        assert isinstance(cut_off, tuple), 'cut_off must be of type int, ' \
+            f'float or tuple, but got {type(cut_off)} instead.'
+    # Auto adjusts contrast for each channel independently and then stacks
+    # the result.
+    s1 = _auto_contrast_channel(img, 0, cut_off)
+    s2 = _auto_contrast_channel(img, 1, cut_off)
+    s3 = _auto_contrast_channel(img, 2, cut_off)
+    contrasted_img = np.stack([s1, s2, s3], axis=-1)
+    return contrasted_img.astype(img.dtype)
+
+
 def adjust_sharpness(img, factor=1., kernel=None):
     """Adjust image sharpness.
 

tests/test_image/test_photometric.py

@@ -202,6 +202,50 @@ def _adjust_contrast(img, factor):
                 rtol=0,
                 atol=1)
 
+    def test_auto_contrast(self, nb_rand_test=100):
+
+        def _auto_contrast(img, cut_off=0):
+            from PIL.ImageOps import autocontrast
+            from PIL import Image
+            # Image.fromarray defaultly supports RGB, not BGR.
+            # convert from BGR to RGB
+            img = Image.fromarray(img[..., ::-1], mode='RGB')
+            contrasted_img = autocontrast(img, cut_off)
+            # convert from RGB to BGR
+            return np.asarray(contrasted_img)[..., ::-1]
+
+        img = np.array([[0, 128, 255], [1, 127, 254], [2, 129, 253]],
+                       dtype=np.uint8)
+        img = np.stack([img, img, img], axis=-1)
+
+        # test case without cut-off
+        assert_array_equal(mmcv.auto_contrast(img), _auto_contrast(img))
+        # test case with cut-off as int
+        assert_array_equal(
+            mmcv.auto_contrast(img, 10), _auto_contrast(img, 10))
+        # test case with cut-off as float
+        assert_array_equal(
+            mmcv.auto_contrast(img, 12.5), _auto_contrast(img, 12.5))
+        # test case with cut-off as tuple
+        assert_array_equal(
+            mmcv.auto_contrast(img, (10, 10)), _auto_contrast(img, 10))
+        # test case with cut-off with sum over 100
+        assert_array_equal(
+            mmcv.auto_contrast(img, 60), _auto_contrast(img, 60))
+
+        # test auto_contrast with randomly sampled images and factors.
+        for _ in range(nb_rand_test):
+            img = np.clip(
+                np.random.uniform(0, 1, (1200, 1000, 3)) * 260, 0,
+                255).astype(np.uint8)
+            # cut-offs are not set as tuple since in `build.yml`, pillow 6.2.2
+            # is installed, which does not support setting low cut-off and high
+            #  cut-off differently.
+            # With pillow above 8.0.0, cut_off can be set as tuple
+            cut_off = np.random.rand() * 100
+            assert_array_equal(
+                mmcv.auto_contrast(img, cut_off), _auto_contrast(img, cut_off))
+
     def test_adjust_sharpness(self, nb_rand_test=100):
 
         def _adjust_sharpness(img, factor):