Add colorspace functions: rgb2ycbcr, bgr2ycbcr, ycbcr2rgb and ycbcr2bgr (#259)

* update

* add functions

* update tests

* resolve comments

* resolve comments

* resolve comments

* update docstr

mmcv/image/__init__.py

 # Copyright (c) Open-MMLab. All rights reserved.
-from .colorspace import (bgr2gray, bgr2hls, bgr2hsv, bgr2rgb, gray2bgr,
-                         gray2rgb, hls2bgr, hsv2bgr, imconvert, rgb2bgr,
-                         rgb2gray)
+from .colorspace import (bgr2gray, bgr2hls, bgr2hsv, bgr2rgb, bgr2ycbcr,
+                         gray2bgr, gray2rgb, hls2bgr, hsv2bgr, imconvert,
+                         rgb2bgr, rgb2gray, rgb2ycbcr, ycbcr2bgr, ycbcr2rgb)
 from .geometric import (imcrop, imflip, imflip_, impad, impad_to_multiple,
                         imrescale, imresize, imresize_like, imrotate,
                         rescale_size)
@@ -15,5 +15,6 @@ __all__ = [
     'imresize', 'imresize_like', 'rescale_size', 'imcrop', 'imflip', 'imflip_',
     'impad', 'impad_to_multiple', 'imrotate', 'imfrombytes', 'imread',
     'imwrite', 'supported_backends', 'use_backend', 'imdenormalize',
-    'imnormalize', 'imnormalize_', 'iminvert', 'posterize', 'solarize'
+    'imnormalize', 'imnormalize_', 'iminvert', 'posterize', 'solarize',
+    'rgb2ycbcr', 'bgr2ycbcr', 'ycbcr2rgb', 'ycbcr2bgr'
 ]

mmcv/image/colorspace.py

 # Copyright (c) Open-MMLab. All rights reserved.
 import cv2
+import numpy as np
 
 
 def imconvert(img, src, dst):
@@ -73,13 +74,204 @@ def gray2rgb(img):
         img (ndarray): The input image.
 
     Returns:
-        ndarray: The converted BGR image.
+        ndarray: The converted RGB image.
     """
     img = img[..., None] if img.ndim == 2 else img
     out_img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
     return out_img
 
 
+def _convert_input_type_range(img):
+    """Convert the type and range of the input image.
+
+    It converts the input image to np.float32 type and range of [0, 1].
+    It is mainly used for pre-processing the input image in colorspace
+    convertion functions such as rgb2ycbcr and ycbcr2rgb.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        (ndarray): The converted image with type of np.float32 and range of
+            [0, 1].
+    """
+    img_type = img.dtype
+    img = img.astype(np.float32)
+    if img_type == np.float32:
+        pass
+    elif img_type == np.uint8:
+        img /= 255.
+    else:
+        raise TypeError('The img type should be np.float32 or np.uint8, '
+                        f'but got {img_type}')
+    return img
+
+
+def _convert_output_type_range(img, dst_type):
+    """Convert the type and range of the image according to dst_type.
+
+    It converts the image to desired type and range. If `dst_type` is np.uint8,
+    images will be converted to np.uint8 type with range [0, 255]. If
+    `dst_type` is np.float32, it converts the image to np.float32 type with
+    range [0, 1].
+    It is mainly used for post-processing images in colorspace convertion
+    functions such as rgb2ycbcr and ycbcr2rgb.
+
+    Args:
+        img (ndarray): The image to be converted with np.float32 type and
+            range [0, 255].
+        dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
+            converts the image to np.uint8 type with range [0, 255]. If
+            dst_type is np.float32, it converts the image to np.float32 type
+            with range [0, 1].
+
+    Returns:
+        (ndarray): The converted image with desired type and range.
+    """
+    if dst_type not in (np.uint8, np.float32):
+        raise TypeError('The dst_type should be np.float32 or np.uint8, '
+                        f'but got {dst_type}')
+    if dst_type == np.uint8:
+        img = img.round()
+    else:
+        img /= 255.
+    return img.astype(dst_type)
+
+
+def rgb2ycbcr(img, y_only=False):
+    """Convert a RGB image to YCbCr image.
+
+    This function produces the same results as Matlab's `rgb2ycbcr` function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
+                  [24.966, 112.0, -18.214]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def bgr2ycbcr(img, y_only=False):
+    """Convert a BGR image to YCbCr image.
+
+    The bgr version of rgb2ycbcr.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
+                  [65.481, -37.797, 112.0]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def ycbcr2rgb(img):
+    """Convert a YCbCr image to RGB image.
+
+    This function produces the same results as Matlab's ycbcr2rgb function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> RGB`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        ndarray: The converted RGB image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621],
+                              [0, -0.00153632, 0.00791071],
+                              [0.00625893, -0.00318811, 0]]) * 255.0 + [
+                                  -222.921, 135.576, -276.836
+                              ]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def ycbcr2bgr(img):
+    """Convert a YCbCr image to BGR image.
+
+    The bgr version of ycbcr2rgb.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> BGR`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        ndarray: The converted BGR image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621],
+                              [0.00791071, -0.00153632, 0],
+                              [0, -0.00318811, 0.00625893]]) * 255.0 + [
+                                  -276.836, 135.576, -222.921
+                              ]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
 def convert_color_factory(src, dst):
 
     code = getattr(cv2, f'COLOR_{src.upper()}2{dst.upper()}')

tests/test_image/test_colorspace.py

@@ -5,6 +5,8 @@ import pytest
 from numpy.testing import assert_array_almost_equal, assert_array_equal
 
 import mmcv
+from mmcv.image.colorspace import (_convert_input_type_range,
+                                   _convert_output_type_range)
 
 
 def test_bgr2gray():
@@ -69,12 +71,10 @@ def test_bgr2hsv():
     in_img = np.random.rand(10, 10, 3).astype(np.float32)
     out_img = mmcv.bgr2hsv(in_img)
     argmax = in_img.argmax(axis=2)
-    computed_hsv = np.empty_like(in_img, dtype=in_img.dtype)
+    computed_hsv = np.empty_like(in_img)
     for i in range(in_img.shape[0]):
         for j in range(in_img.shape[1]):
-            b = in_img[i, j, 0]
-            g = in_img[i, j, 1]
-            r = in_img[i, j, 2]
+            b, g, r = in_img[i, j]
             v = max(r, g, b)
             s = (v - min(r, g, b)) / v if v != 0 else 0
             if argmax[i, j] == 0:
@@ -89,16 +89,231 @@ def test_bgr2hsv():
     assert_array_almost_equal(out_img, computed_hsv, decimal=2)
 
 
+def test_convert_input_type_range():
+    with pytest.raises(TypeError):
+        # The img type should be np.float32 or np.uint8
+        in_img = np.random.rand(10, 10, 3).astype(np.uint64)
+        _convert_input_type_range(in_img)
+    # np.float32
+    in_img = np.random.rand(10, 10, 3).astype(np.float32)
+    out_img = _convert_input_type_range(in_img)
+    assert out_img.dtype == np.float32
+    assert np.absolute(out_img).mean() < 1
+    # np.uint8
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = _convert_input_type_range(in_img)
+    assert out_img.dtype == np.float32
+    assert np.absolute(out_img).mean() < 1
+
+
+def test_convert_output_type_range():
+    with pytest.raises(TypeError):
+        # The dst_type should be np.float32 or np.uint8
+        in_img = np.random.rand(10, 10, 3).astype(np.float32)
+        _convert_output_type_range(in_img, np.uint64)
+    # np.float32
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.float32)
+    out_img = _convert_output_type_range(in_img, np.float32)
+    assert out_img.dtype == np.float32
+    assert np.absolute(out_img).mean() < 1
+    # np.uint8
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.float32)
+    out_img = _convert_output_type_range(in_img, np.uint8)
+    assert out_img.dtype == np.uint8
+    assert np.absolute(out_img).mean() > 1
+
+
+def test_rgb2ycbcr():
+    with pytest.raises(TypeError):
+        # The img type should be np.float32 or np.uint8
+        in_img = np.random.rand(10, 10, 3).astype(np.uint64)
+        mmcv.rgb2ycbcr(in_img)
+
+    # float32
+    in_img = np.random.rand(10, 10, 3).astype(np.float32)
+    out_img = mmcv.rgb2ycbcr(in_img)
+    computed_ycbcr = np.empty_like(in_img)
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            r, g, b = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            cb = 128 - r * 37.797 - g * 74.203 + b * 112.0
+            cr = 128 + r * 112.0 - g * 93.786 - b * 18.214
+            computed_ycbcr[i, j, :] = [y, cb, cr]
+    computed_ycbcr /= 255.
+    assert_array_almost_equal(out_img, computed_ycbcr, decimal=2)
+    # y_only=True
+    out_img = mmcv.rgb2ycbcr(in_img, y_only=True)
+    computed_y = np.empty_like(out_img, dtype=out_img.dtype)
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            r, g, b = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            computed_y[i, j] = y
+    computed_y /= 255.
+    assert_array_almost_equal(out_img, computed_y, decimal=2)
+
+    # uint8
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = mmcv.rgb2ycbcr(in_img)
+    computed_ycbcr = np.empty_like(in_img)
+    in_img = in_img / 255.
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            r, g, b = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            cb = 128 - r * 37.797 - g * 74.203 + b * 112.0
+            cr = 128 + r * 112.0 - g * 93.786 - b * 18.214
+            y, cb, cr = y.round(), cb.round(), cr.round()
+            computed_ycbcr[i, j, :] = [y, cb, cr]
+    assert_array_almost_equal(out_img, computed_ycbcr, decimal=2)
+    # y_only=True
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = mmcv.rgb2ycbcr(in_img, y_only=True)
+    computed_y = np.empty_like(out_img, dtype=out_img.dtype)
+    in_img = in_img / 255.
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            r, g, b = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            y = y.round()
+            computed_y[i, j] = y
+    assert_array_almost_equal(out_img, computed_y, decimal=2)
+
+
+def test_bgr2ycbcr():
+    # float32
+    in_img = np.random.rand(10, 10, 3).astype(np.float32)
+    out_img = mmcv.bgr2ycbcr(in_img)
+    computed_ycbcr = np.empty_like(in_img)
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            b, g, r = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            cb = 128 - r * 37.797 - g * 74.203 + b * 112.0
+            cr = 128 + r * 112.0 - g * 93.786 - b * 18.214
+            computed_ycbcr[i, j, :] = [y, cb, cr]
+    computed_ycbcr /= 255.
+    assert_array_almost_equal(out_img, computed_ycbcr, decimal=2)
+    # y_only=True
+    in_img = np.random.rand(10, 10, 3).astype(np.float32)
+    out_img = mmcv.bgr2ycbcr(in_img, y_only=True)
+    computed_y = np.empty_like(out_img, dtype=out_img.dtype)
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            b, g, r = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            computed_y[i, j] = y
+    computed_y /= 255.
+    assert_array_almost_equal(out_img, computed_y, decimal=2)
+
+    # uint8
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = mmcv.bgr2ycbcr(in_img)
+    computed_ycbcr = np.empty_like(in_img)
+    in_img = in_img / 255.
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            b, g, r = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            cb = 128 - r * 37.797 - g * 74.203 + b * 112.0
+            cr = 128 + r * 112.0 - g * 93.786 - b * 18.214
+            y, cb, cr = y.round(), cb.round(), cr.round()
+            computed_ycbcr[i, j, :] = [y, cb, cr]
+    assert_array_almost_equal(out_img, computed_ycbcr, decimal=2)
+    # y_only = True
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = mmcv.bgr2ycbcr(in_img, y_only=True)
+    computed_y = np.empty_like(out_img, dtype=out_img.dtype)
+    in_img = in_img / 255.
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            b, g, r = in_img[i, j]
+            y = 16 + r * 65.481 + g * 128.553 + b * 24.966
+            y = y.round()
+            computed_y[i, j] = y
+    assert_array_almost_equal(out_img, computed_y, decimal=2)
+
+
+def test_ycbcr2rgb():
+    with pytest.raises(TypeError):
+        # The img type should be np.float32 or np.uint8
+        in_img = np.random.rand(10, 10, 3).astype(np.uint64)
+        mmcv.ycbcr2rgb(in_img)
+
+    # float32
+    in_img = np.random.rand(10, 10, 3).astype(np.float32)
+    out_img = mmcv.ycbcr2rgb(in_img)
+    computed_rgb = np.empty_like(in_img)
+    in_img *= 255.
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            y, cb, cr = in_img[i, j]
+            r = -222.921 + y * 0.00456621 * 255 + cr * 0.00625893 * 255
+            g = 135.576 + y * 0.00456621 * 255 - cb * 0.00153632 * 255 - \
+                cr * 0.00318811 * 255
+            b = -276.836 + y * 0.00456621 * 255. + cb * 0.00791071 * 255
+            computed_rgb[i, j, :] = [r, g, b]
+    computed_rgb /= 255.
+    assert_array_almost_equal(out_img, computed_rgb, decimal=2)
+
+    # uint8
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = mmcv.ycbcr2rgb(in_img)
+    computed_rgb = np.empty_like(in_img)
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            y, cb, cr = in_img[i, j]
+            r = -222.921 + y * 0.00456621 * 255 + cr * 0.00625893 * 255
+            g = 135.576 + y * 0.00456621 * 255 - cb * 0.00153632 * 255 - \
+                cr * 0.00318811 * 255
+            b = -276.836 + y * 0.00456621 * 255. + cb * 0.00791071 * 255
+            r, g, b = r.round(), g.round(), b.round()
+            computed_rgb[i, j, :] = [r, g, b]
+    assert_array_almost_equal(out_img, computed_rgb, decimal=2)
+
+
+def test_ycbcr2bgr():
+    # float32
+    in_img = np.random.rand(10, 10, 3).astype(np.float32)
+    out_img = mmcv.ycbcr2bgr(in_img)
+    computed_bgr = np.empty_like(in_img)
+    in_img *= 255.
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            y, cb, cr = in_img[i, j]
+            r = -222.921 + y * 0.00456621 * 255 + cr * 0.00625893 * 255
+            g = 135.576 + y * 0.00456621 * 255 - cb * 0.00153632 * 255 - \
+                cr * 0.00318811 * 255
+            b = -276.836 + y * 0.00456621 * 255. + cb * 0.00791071 * 255
+            computed_bgr[i, j, :] = [b, g, r]
+    computed_bgr /= 255.
+    assert_array_almost_equal(out_img, computed_bgr, decimal=2)
+
+    # uint8
+    in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
+    out_img = mmcv.ycbcr2bgr(in_img)
+    computed_bgr = np.empty_like(in_img)
+    for i in range(in_img.shape[0]):
+        for j in range(in_img.shape[1]):
+            y, cb, cr = in_img[i, j]
+            r = -222.921 + y * 0.00456621 * 255 + cr * 0.00625893 * 255
+            g = 135.576 + y * 0.00456621 * 255 - cb * 0.00153632 * 255 - \
+                cr * 0.00318811 * 255
+            b = -276.836 + y * 0.00456621 * 255. + cb * 0.00791071 * 255
+            r, g, b = r.round(), g.round(), b.round()
+            computed_bgr[i, j, :] = [b, g, r]
+    assert_array_almost_equal(out_img, computed_bgr, decimal=2)
+
+
 def test_bgr2hls():
     in_img = np.random.rand(10, 10, 3).astype(np.float32)
     out_img = mmcv.bgr2hls(in_img)
     argmax = in_img.argmax(axis=2)
-    computed_hls = np.empty_like(in_img, dtype=in_img.dtype)
+    computed_hls = np.empty_like(in_img)
     for i in range(in_img.shape[0]):
         for j in range(in_img.shape[1]):
-            b = in_img[i, j, 0]
-            g = in_img[i, j, 1]
-            r = in_img[i, j, 2]
+            b, g, r = in_img[i, j]
             maxc = max(r, g, b)
             minc = min(r, g, b)
             _l = (minc + maxc) / 2.0