|
- # Copyright (c) Open-MMLab. All rights reserved.
- import cv2
- import numpy as np
-
-
- def imconvert(img, src, dst):
- """Convert an image from the src colorspace to dst colorspace.
-
- Args:
- img (ndarray): The input image.
- src (str): The source colorspace, e.g., 'rgb', 'hsv'.
- dst (str): The destination colorspace, e.g., 'rgb', 'hsv'.
-
- Returns:
- ndarray: The converted image.
- """
- code = getattr(cv2, f'COLOR_{src.upper()}2{dst.upper()}')
- out_img = cv2.cvtColor(img, code)
- return out_img
-
-
- def bgr2gray(img, keepdim=False):
- """Convert a BGR image to grayscale image.
-
- Args:
- img (ndarray): The input image.
- keepdim (bool): If False (by default), then return the grayscale image
- with 2 dims, otherwise 3 dims.
-
- Returns:
- ndarray: The converted grayscale image.
- """
- out_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
- if keepdim:
- out_img = out_img[..., None]
- return out_img
-
-
- def rgb2gray(img, keepdim=False):
- """Convert a RGB image to grayscale image.
-
- Args:
- img (ndarray): The input image.
- keepdim (bool): If False (by default), then return the grayscale image
- with 2 dims, otherwise 3 dims.
-
- Returns:
- ndarray: The converted grayscale image.
- """
- out_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
- if keepdim:
- out_img = out_img[..., None]
- return out_img
-
-
- def gray2bgr(img):
- """Convert a grayscale image to BGR image.
-
- Args:
- img (ndarray): The input image.
-
- Returns:
- ndarray: The converted BGR image.
- """
- img = img[..., None] if img.ndim == 2 else img
- out_img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
- return out_img
-
-
- def gray2rgb(img):
- """Convert a grayscale image to RGB image.
-
- Args:
- img (ndarray): The input image.
-
- Returns:
- 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()}')
-
- def convert_color(img):
- out_img = cv2.cvtColor(img, code)
- return out_img
-
- convert_color.__doc__ = f"""Convert a {src.upper()} image to {dst.upper()}
- image.
-
- Args:
- img (ndarray or str): The input image.
-
- Returns:
- ndarray: The converted {dst.upper()} image.
- """
-
- return convert_color
-
-
- bgr2rgb = convert_color_factory('bgr', 'rgb')
-
- rgb2bgr = convert_color_factory('rgb', 'bgr')
-
- bgr2hsv = convert_color_factory('bgr', 'hsv')
-
- hsv2bgr = convert_color_factory('hsv', 'bgr')
-
- bgr2hls = convert_color_factory('bgr', 'hls')
-
- hls2bgr = convert_color_factory('hls', 'bgr')
|