[Feature] Add LUT transform (#646)

* add lut_transform

* restore

* fix test bug

* add test

* Fix syntax bug

* Delete test id

* modfiy docstrinng

* simplify assert

mmcv/image/__init__.py

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

mmcv/image/photometric.py

@@ -224,3 +224,27 @@ def adjust_contrast(img, factor=1.):
         img.astype(np.float32), factor, degenerated.astype(np.float32),
         1 - factor, 0)
     return contrasted_img.astype(img.dtype)
+
+
+def lut_transform(img, lut_table):
+    """Transform array by look-up table.
+
+    The function lut_transform fills the output array with values from the
+    look-up table. Indices of the entries are taken from the input array.
+
+    Args:
+        img (ndarray): Image to be transformed.
+        lut_table (ndarray): look-up table of 256 elements; in case of
+            multi-channel input array, the table should either have a single
+            channel (in this case the same table is used for all channels) or
+            the same number of channels as in the input array.
+
+    Returns:
+        ndarray: The transformed image.
+    """
+    assert isinstance(img, np.ndarray)
+    assert 0 <= np.min(img) and np.max(img) <= 255
+    assert isinstance(lut_table, np.ndarray)
+    assert lut_table.shape == (256, )
+
+    return cv2.LUT(np.array(img, dtype=np.uint8), lut_table)

tests/test_image/test_photometric.py

@@ -200,3 +200,22 @@ class TestPhotometric:
                 _adjust_contrast(img, factor).astype(np.int32),
                 rtol=0,
                 atol=1)
+
+    def test_lut_transform(self):
+        lut_table = np.array(list(range(256)))
+
+        img = mmcv.lut_transform(self.img, lut_table)
+        baseline = cv2.LUT(self.img, lut_table)
+        assert np.allclose(img, baseline)
+
+        input_img = np.array(
+            [[[0, 128, 255], [255, 128, 0]], [[0, 128, 255], [255, 128, 0]]],
+            dtype=np.float)
+        img = mmcv.lut_transform(input_img, lut_table)
+        baseline = cv2.LUT(np.array(input_img, dtype=np.uint8), lut_table)
+        assert np.allclose(img, baseline)
+
+        input_img = np.random.randint(0, 256, size=(7, 8, 9, 10, 11))
+        img = mmcv.lut_transform(input_img, lut_table)
+        baseline = cv2.LUT(np.array(input_img, dtype=np.uint8), lut_table)
+        assert np.allclose(img, baseline)