Supports brightness and contrast augmentations (#546)

* add brightness and contrast augmentation

* remove unnecessary

* reformat

* relax the precision constrain for adjust_brightness aug

* fix percision assertion error in unit test

* remove toy

* rename alpha as factor

* use np.testing.assert_allclose in place of np.less_equal

mmcv/image/__init__.py

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

mmcv/image/photometric.py

@@ -1,7 +1,7 @@
 import cv2
 import numpy as np
 
-from .colorspace import bgr2gray
+from .colorspace import bgr2gray, gray2bgr
 
 
 def imnormalize(img, mean, std, to_rgb=True):
@@ -166,3 +166,61 @@ def _scale_channel(im, c):
     s3 = _scale_channel(img, 2)
     equalized_img = np.stack([s1, s2, s3], axis=-1)
     return equalized_img
+
+
+def adjust_brightness(img, factor=1.):
+    """Adjust image brightness.
+
+    This function controls the brightness of an image. An
+    enhancement factor of 0.0 gives a black image.
+    A factor of 1.0 gives the original image. This function
+    blends the source image and the degenerated black image:
+
+    ``output = img * factor + degenerated * (1 - factor)``
+
+    Args:
+        img (ndarray): Image to be brightened.
+        factor (float): A value controls the enhancement.
+            Factor 1.0 returns the original image, lower
+            factors mean less color (brightness, contrast,
+            etc), and higher values more. Default 1.
+
+    Returns:
+        ndarray: The brightened image.
+    """
+    degenerated = np.zeros_like(img)
+    # Note manually convert the dtype to np.float32, to
+    # achieve as close results as PIL.ImageEnhance.Brightness.
+    # Set beta=1-factor, and gamma=0
+    brightened_img = cv2.addWeighted(
+        img.astype(np.float32), factor, degenerated.astype(np.float32),
+        1 - factor, 0)
+    return brightened_img.astype(img.dtype)
+
+
+def adjust_contrast(img, factor=1.):
+    """Adjust image contrast.
+
+    This function controls the contrast of an image. An
+    enhancement factor of 0.0 gives a solid grey
+    image. A factor of 1.0 gives the original image. It
+    blends the source image and the degenerated mean image:
+
+    ``output = img * factor + degenerated * (1 - factor)``
+
+    Args:
+        img (ndarray): Image to be contrasted. BGR order.
+        factor (float): Same as :func:`mmcv.adjust_brightness`.
+
+    Returns:
+        ndarray: The contrasted image.
+    """
+    gray_img = bgr2gray(img)
+    hist = np.histogram(gray_img, 256, (0, 255))[0]
+    mean = round(np.sum(gray_img) / np.sum(hist))
+    degenerated = (np.ones_like(img[..., 0]) * mean).astype(img.dtype)
+    degenerated = gray2bgr(degenerated)
+    contrasted_img = cv2.addWeighted(
+        img.astype(np.float32), factor, degenerated.astype(np.float32),
+        1 - factor, 0)
+    return contrasted_img.astype(img.dtype)

tests/test_image/test_photometric.py

@@ -135,3 +135,68 @@ def _imequalize(img):
                 255).astype(np.uint8)
             equalized_img = mmcv.imequalize(img)
             assert_array_equal(equalized_img, _imequalize(img))
+
+    def test_adjust_brightness(self, nb_rand_test=100):
+
+        def _adjust_brightness(img, factor):
+            # adjust the brightness of image using
+            # PIL.ImageEnhance.Brightness
+            from PIL.ImageEnhance import Brightness
+            from PIL import Image
+            img = Image.fromarray(img)
+            brightened_img = Brightness(img).enhance(factor)
+            return np.asarray(brightened_img)
+
+        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 with factor 1.0
+        assert_array_equal(mmcv.adjust_brightness(img, 1.), img)
+        # test case with factor 0.0
+        assert_array_equal(mmcv.adjust_brightness(img, 0.), np.zeros_like(img))
+        # test adjust_brightness with randomly sampled images and factors.
+        for _ in range(nb_rand_test):
+            img = np.clip(
+                np.random.uniform(0, 1, (1000, 1200, 3)) * 260, 0,
+                255).astype(np.uint8)
+            factor = np.random.uniform()
+            np.testing.assert_allclose(
+                mmcv.adjust_brightness(img, factor).astype(np.int32),
+                _adjust_brightness(img, factor).astype(np.int32),
+                rtol=0,
+                atol=1)
+
+    def test_adjust_contrast(self, nb_rand_test=100):
+
+        def _adjust_contrast(img, factor):
+            from PIL.ImageEnhance import Contrast
+            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 = Contrast(img).enhance(factor)
+            # 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 with factor 1.0
+        assert_array_equal(mmcv.adjust_contrast(img, 1.), img)
+        # test case with factor 0.0
+        assert_array_equal(
+            mmcv.adjust_contrast(img, 0.), _adjust_contrast(img, 0.))
+        # test adjust_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)
+            factor = np.random.uniform()
+            # Note the gap (less_equal 1) between PIL.ImageEnhance.Contrast
+            # and mmcv.adjust_contrast comes from the gap that converts from
+            # a color image to gray image using mmcv or PIL.
+            np.testing.assert_allclose(
+                mmcv.adjust_contrast(img, factor).astype(np.int32),
+                _adjust_contrast(img, factor).astype(np.int32),
+                rtol=0,
+                atol=1)