Added solarize and posterize image transformations (#132)

* Added solarize and posterize image transformations

* Added solarize and posterize transformations. Modified __init__.py

* Updated mmcv/image/__init__.py

* Updated the docstrings for solarize and posterize

* add unit tests

* fix a test case

mmcv/image/__init__.py

 from .io import imread, imwrite, imfrombytes
-from .transforms import (bgr2gray, gray2bgr, bgr2rgb, rgb2bgr, bgr2hsv,
-                         hsv2bgr, bgr2hls, hls2bgr, iminvert, imflip, imrotate,
-                         imcrop, impad, impad_to_multiple, imnormalize,
-                         imdenormalize, imresize, imresize_like, imrescale)
+from .transforms import (solarize, posterize, bgr2gray, gray2bgr, bgr2rgb,
+                         rgb2bgr, bgr2hsv, hsv2bgr, bgr2hls, hls2bgr, iminvert,
+                         imflip, imrotate, imcrop, impad, impad_to_multiple,
+                         imnormalize, imdenormalize, imresize, imresize_like,
+                         imrescale)
 
 __all__ = [
-    'imread', 'imwrite', 'imfrombytes', 'bgr2gray', 'gray2bgr', 'bgr2rgb',
-    'rgb2bgr', 'bgr2hsv', 'hsv2bgr', 'bgr2hls', 'hls2bgr', 'iminvert',
-    'imflip', 'imrotate', 'imcrop', 'impad', 'impad_to_multiple',
-    'imnormalize', 'imdenormalize', 'imresize', 'imresize_like', 'imrescale'
+    'solarize', 'posterize', 'imread', 'imwrite', 'imfrombytes', 'bgr2gray',
+    'gray2bgr', 'bgr2rgb', 'rgb2bgr', 'bgr2hsv', 'hsv2bgr', 'bgr2hls',
+    'hls2bgr', 'iminvert', 'imflip', 'imrotate', 'imcrop', 'impad',
+    'impad_to_multiple', 'imnormalize', 'imdenormalize', 'imresize',
+    'imresize_like', 'imrescale'
 ]

mmcv/image/transforms/__init__.py

-from .colorspace import (bgr2gray, gray2bgr, bgr2rgb, rgb2bgr, bgr2hsv,
-                         hsv2bgr, bgr2hls, hls2bgr, iminvert)
+from .colorspace import (solarize, posterize, bgr2gray, gray2bgr, bgr2rgb,
+                         rgb2bgr, bgr2hsv, hsv2bgr, bgr2hls, hls2bgr, iminvert)
 from .geometry import imflip, imrotate, imcrop, impad, impad_to_multiple
 from .normalize import imnormalize, imdenormalize
 from .resize import imresize, imresize_like, imrescale
 
 __all__ = [
-    'bgr2gray', 'gray2bgr', 'bgr2rgb', 'rgb2bgr', 'bgr2hsv', 'hsv2bgr',
-    'bgr2hls', 'hls2bgr', 'iminvert', 'imflip', 'imrotate', 'imcrop', 'impad',
-    'impad_to_multiple', 'imnormalize', 'imdenormalize', 'imresize',
-    'imresize_like', 'imrescale'
+    'solarize', 'posterize', 'bgr2gray', 'gray2bgr', 'bgr2rgb', 'rgb2bgr',
+    'bgr2hsv', 'hsv2bgr', 'bgr2hls', 'hls2bgr', 'iminvert', 'imflip',
+    'imrotate', 'imcrop', 'impad', 'impad_to_multiple', 'imnormalize',
+    'imdenormalize', 'imresize', 'imresize_like', 'imrescale'
 ]

mmcv/image/transforms/colorspace.py

@@ -2,6 +2,35 @@ import cv2
 import numpy as np
 
 
+def solarize(img, thr=128):
+    """Solarize an image (invert all pixel values above a threshold)
+
+    Args:
+        img (ndarray): Image to be solarized.
+        thr (int): Threshold for solarizing (0 - 255).
+
+    Returns:
+        ndarray: The solarized image.
+    """
+    img = np.where(img < thr, img, 255 - img)
+    return img
+
+
+def posterize(img, bits):
+    """Posterize an image (reduce the number of bits for each color channel)
+
+    Args:
+        img (ndarray): Image to be posterized.
+        bits (int): Number of bits (1 to 8) to use for posterizing.
+
+    Returns:
+        ndarray: The posterized image.
+    """
+    shift = 8 - bits
+    img = np.left_shift(np.right_shift(img, shift), shift)
+    return img
+
+
 def iminvert(img):
     """Invert (negate) an image
     Args:

tests/test_image.py

@@ -325,3 +325,23 @@ class TestImage(object):
         img_r = np.array([[255, 127, 0], [254, 128, 1], [253, 126, 2]],
                          dtype=np.uint8)
         assert_array_equal(mmcv.iminvert(img), img_r)
+
+    def test_solarize(self):
+        img = np.array([[0, 128, 255], [1, 127, 254], [2, 129, 253]],
+                       dtype=np.uint8)
+        img_r = np.array([[0, 127, 0], [1, 127, 1], [2, 126, 2]],
+                         dtype=np.uint8)
+        assert_array_equal(mmcv.solarize(img), img_r)
+        img_r = np.array([[0, 127, 0], [1, 128, 1], [2, 126, 2]],
+                         dtype=np.uint8)
+        assert_array_equal(mmcv.solarize(img, 100), img_r)
+
+    def test_posterize(self):
+        img = np.array([[0, 128, 255], [1, 127, 254], [2, 129, 253]],
+                       dtype=np.uint8)
+        img_r = np.array([[0, 128, 128], [0, 0, 128], [0, 128, 128]],
+                         dtype=np.uint8)
+        assert_array_equal(mmcv.posterize(img, 1), img_r)
+        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)