颜色空间转换

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

概述

在计算机视觉中，所选颜色空间可能会对模型的性能产生重大影响。虽然 RGB 是最常见的颜色空间，但在许多情况下，模型在切换到其他颜色空间（如 YUV、YCbCr、XYZ (CIE) 等）时表现更好。

tensorflow-io 包提供了一系列颜色空间转换 API，可用于准备和增强图像数据。

设置

安装所需的软件包，并重新启动运行时

pip install -q tensorflow-io

下载示例图像

本教程中使用的图像示例是 雪地中的猫，但可以替换为任何 JPEG 图像。

以下操作将下载图像并将其保存到本地磁盘，文件名为 sample.jpg

curl -o sample.jpg -L https://storage.googleapis.com/download.tensorflow.org/example_images/320px-Felis_catus-cat_on_snow.jpg

ls -ls sample.jpg

% Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 17858  100 17858    0     0   235k      0 --:--:-- --:--:-- --:--:--  235k
20 -rw-rw-r-- 1 kbuilder kokoro 17858 Oct 27 16:33 sample.jpg

用法

读取图像文件

读取并解码图像，将其转换为形状为 (213, 320, 3) 的 uint8 张量

import tensorflow as tf
import tensorflow_io as tfio

image = tf.image.decode_jpeg(tf.io.read_file('sample.jpg'))

print(image.shape, image.dtype)

(213, 320, 3) <dtype: 'uint8'>

可以通过以下方式显示图像

import matplotlib.pyplot as plt

plt.figure()
plt.imshow(image)
plt.axis('off')
plt.show()

png

将 RGB 转换为灰度

可以使用 tfio.experimental.color.rgb_to_grayscale 将 RGB 图像转换为 灰度，从而将通道从 3 减少到 1。

grayscale = tfio.experimental.color.rgb_to_grayscale(image)

print(grayscale.shape, grayscale.dtype)

# use tf.squeeze to remove last channel for plt.imshow to display:
plt.figure()
plt.imshow(tf.squeeze(grayscale, axis=-1), cmap='gray')
plt.axis('off')
plt.show()

(213, 320, 1) <dtype: 'uint8'>

png

将 RGB 转换为 BGR

一些图像软件和相机制造商可能更喜欢 BGR，可以通过 tfio.experimental.color.rgb_to_bgr 获取。

bgr = tfio.experimental.color.rgb_to_bgr(image)

print(bgr.shape, bgr.dtype)

plt.figure()
plt.imshow(bgr)
plt.axis('off')
plt.show()

(213, 320, 3) <dtype: 'uint8'>

png

将 RGB 转换为 CIE XYZ

CIE XYZ（或 CIE 1931 XYZ）是许多图像处理程序中使用的常见颜色空间。以下是通过 tfio.experimental.color.rgb_to_xyz 从 RGB 转换为 CIE XYZ 的转换。请注意，tfio.experimental.color.rgb_to_xyz 假设浮点输入范围为 [0, 1]，因此需要额外的预处理。

# convert to float32
image_float32 = tf.cast(image, tf.float32) / 255.0

xyz_float32 = tfio.experimental.color.rgb_to_xyz(image_float32)

# convert back uint8
xyz = tf.cast(xyz_float32 * 255.0, tf.uint8)

print(xyz.shape, xyz.dtype)

plt.figure()
plt.imshow(xyz)
plt.axis('off')
plt.show()

(213, 320, 3) <dtype: 'uint8'>

png

将 RGB 转换为 YCbCr

最后，YCbCr 是许多视频系统中的默认颜色空间。可以通过 tfio.experimental.color.rgb_to_ycbcr 转换为 YCbCr。

ycbcr = tfio.experimental.color.rgb_to_ycbcr(image)

print(ycbcr.shape, ycbcr.dtype)

plt.figure()
plt.imshow(ycbcr, cmap='gray')
plt.axis('off')
plt.show()

(213, 320, 3) <dtype: 'uint8'>

png

更有趣的是，YCbCr 可以分解为 Y'（亮度）、Cb（蓝色差色度）和 Cr（红色差色度）分量，每个分量都承载着感知上有意义的信息。

y, cb, cr = ycbcr[:,:,0], ycbcr[:,:,1], ycbcr[:,:,2]

# Y' component
plt.figure()
plt.imshow(y, cmap='gray')
plt.axis('off')
plt.show()

# Cb component
plt.figure()
plt.imshow(cb, cmap='gray')
plt.axis('off')
plt.show()

# Cr component
plt.figure()
plt.imshow(cr, cmap='gray')
plt.axis('off')
plt.show()

png

png

png