Supports Color augmentation. (#542)

* add Color augmentation * reformat * reformat * reformat docstring * reformat docstring * add more unit test * add more unit test * add clip value and uint test for image with type float * rename function name
open-mmlab · Sep 8, 2020 · 95417a5 · 95417a5
1 parent fec7cd6
commit 95417a5
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 3 deletions.
diff --git a/mmcv/image/__init__.py b/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 (adjust_color, 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', 'adjust_color'
 ]
diff --git a/mmcv/image/photometric.py b/mmcv/image/photometric.py
@@ -1,6 +1,8 @@
 import cv2
 import numpy as np
 
+from .colorspace import bgr2gray
+
 
 def imnormalize(img, mean, std, to_rgb=True):
     """Normalize an image with mean and std.
@@ -91,3 +93,33 @@ def posterize(img, bits):
     shift = 8 - bits
     img = np.left_shift(np.right_shift(img, shift), shift)
     return img
+
+
+def adjust_color(img, alpha=1, beta=None, gamma=0):
+    """It blends the source image and its gray image:
+
+    ``output = img * alpha + gray_img * beta + gamma``
+
+    Args:
+        img (ndarray): The input source image.
+        alpha (int | float): Weight for the source image. Default 1.
+        beta (int | float): Weight for the converted gray image.
+            If None, it's assigned the value (1 - `alpha`).
+        gamma (int | float): Scalar added to each sum.
+            Same as :func:`cv2.addWeighted`. Default 0.
+
+    Returns:
+        ndarray: Colored image which has the same size and dtype as input.
+    """
+    gray_img = bgr2gray(img)
+    gray_img = np.tile(gray_img[..., None], [1, 1, 3])
+    if beta is None:
+        beta = 1 - alpha
+    colored_img = cv2.addWeighted(img, alpha, gray_img, beta, gamma)
+    if not colored_img.dtype == np.uint8:
+        # Note when the dtype of `img` is not defaultly `np.uint8`
+        # (e.g. np.float32), the value in `colored_img` got from cv2
+        # is not guaranteed to be in range [0, 255], so here clip
+        # is needed.
+        colored_img = np.clip(colored_img, 0, 255)
+    return colored_img
diff --git a/tests/test_image/test_photometric.py b/tests/test_image/test_photometric.py
@@ -75,3 +75,34 @@ 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_adjust_color(self):
+        img = np.array([[0, 128, 255], [1, 127, 254], [2, 129, 253]],
+                       dtype=np.uint8)
+        img = np.stack([img, img, img], axis=-1)
+        assert_array_equal(mmcv.adjust_color(img), img)
+        img_gray = mmcv.bgr2gray(img)
+        img_r = np.stack([img_gray, img_gray, img_gray], axis=-1)
+        assert_array_equal(mmcv.adjust_color(img, 0), img_r)
+        assert_array_equal(mmcv.adjust_color(img, 0, 1), img_r)
+        assert_array_equal(
+            mmcv.adjust_color(img, 0.5, 0.5),
+            np.round(np.clip((img * 0.5 + img_r * 0.5), 0,
+                             255)).astype(img.dtype))
+        assert_array_equal(
+            mmcv.adjust_color(img, 1, 1.5),
+            np.round(np.clip(img * 1 + img_r * 1.5, 0, 255)).astype(img.dtype))
+        assert_array_equal(
+            mmcv.adjust_color(img, 0.8, -0.6, gamma=2),
+            np.round(np.clip(img * 0.8 - 0.6 * img_r + 2, 0,
+                             255)).astype(img.dtype))
+        assert_array_equal(
+            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)).astype(img.dtype))
+
+        # test float type of image
+        img = img.astype(np.float32)
+        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)))