如何使用NumPy将sRGB转换为NV12格式？

Question

NV12格式定义了YUV色彩空间的特定颜色通道顺序，带有420子采样。
NV12格式主要用于视频编解码流程中。
libyuv对NV12的描述：

NV12是一种双平面格式，具有完整大小的Y平面，后跟一个单色度平面，其中编织了U和V值。NV21与之相同，但编织了V和U值。NV12中的12指的是每像素12位。NV12具有半宽度和半高度的色度通道，因此是420子采样。

在NV12的上下文中，YUV格式主要被称为YCbCr颜色空间。
NV12元素为8位每元素（uint8类型）。
在本文中，YUV元素遵循“有限范围”标准：Y范围为[16, 235]，U、V范围为[16, 240]。

sRGB（标准红绿蓝）是PC系统使用的标准色彩空间。
在帖子的上下文中，sRGB颜色分量范围为[0,255]（uint8类型）。
RGB元素排序与帖子无关（假设有3个颜色平面）。

目前至少有两种应用NV12的YCbCr格式：

• BT.601 - 适用于SDTV。
• BT.709 - 适用于HDTV。

NV12元素排序示例：
YYYYYY
YYYYYY
UVUVUV

RGB转NV12的转换可以通过以下步骤描述：

• 色彩空间转换 - 从sRGB转换为YUV色彩空间。
• 色度抽样 - 在每个轴上将U、V通道缩小2倍（从YUV444转换为YUV420）。
• 色度元素交错 - 将U、V元素排列为U、V、U、V...

下图说明了应用6x6像素图像大小的转换步骤：

如何使用NumPy将sRGB转换为NV12？

注意：
该问题涉及Python实现，演示转换过程（本文不旨在使用现有函数，如OpenCV实现）。

- Rotem

Answer 1

使用NumPy将sRGB转换为NV12格式

本文旨在演示转换过程。
下面的Python实现使用NumPy，有意避免使用OpenCV。

RGB到NV12转换步骤：

• 色彩空间转换-从sRGB转换为YUV色彩空间：
  使用sRGB到YCbCr转换公式。
  将每个RGB三元组乘以3x3转换矩阵，并加上一个3个偏移量的向量。
  该文章展示了BT.709和BT.601两种转换（唯一的区别是系数矩阵）。
• 色度下采样-在每个轴上缩小U、V通道的因子x2（从YUV444转换为YUV420）。
  该实现使用双线性插值将U、V调整了0.5倍大小。
  注意：双线性插值不是最优的下采样方法，但通常足够好。
  代码不使用cv2.resize，而是使用每2x2像素的平均值（结果等效于双线性插值）。
  注意：如果输入分辨率两个维度不同，则实现会失败。
• 色度元素交错-将U、V元素排列为U、V、U、V...
  通过数组索引操作实现。

以下是将RGB转换为NV12标准的Python代码示例：

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import subprocess as sp  # The module is used for testing (using FFmpeg as reference).

do_use_bt709 = True  # True for BT.709, False for BT.601

rgb = mpimg.imread('rgb_input.png')*255.0   # Read RGB input image, multiply by 255 (set RGB range to [0, 255]).
r, g, b = np.squeeze(np.split(rgb, 3, -1))  # Split RGB to R, G and B numpy arrays.
rows, cols = r.shape

# I. Convert RGB to YUV (convert sRGB to YUV444)
#################################################
if do_use_bt709:
    # Convert sRGB to YUV, BT.709 standard
    # Conversion formula used: 8 bit sRGB to "limited range" 8 bit YUV (BT.709).            
    y =  0.1826*r + 0.6142*g + 0.0620*b + 16
    u = -0.1006*r - 0.3386*g + 0.4392*b + 128
    v =  0.4392*r - 0.3989*g - 0.0403*b + 128
else:
    # Convert sRGB to YUV, BT.601 standard.
    # Conversion formula used: 8 bit sRGB to "limited range" 8 bit YUV (BT.601).
    y =  0.2568*r + 0.5041*g + 0.0979*b + 16
    u = -0.1482*r - 0.2910*g + 0.4392*b + 128
    v =  0.4392*r - 0.3678*g - 0.0714*b + 128


# II. U,V Downscaling (convert YUV444 to YUV420)
##################################################
# Shrink U and V channels by a factor of x2 in each axis (use bi-linear interpolation).
#shrunk_u = cv2.resize(u, (cols//2, rows//2), interpolation=cv2.INTER_LINEAR)
#shrunk_v = cv2.resize(v, (cols//2, rows//2), interpolation=cv2.INTER_LINEAR)

# Each element of shrunkU is the mean of 2x2 elements of U
# Result is equivalent to resize by a factor of 0.5 with bi-linear interpolation.
shrunk_u = (u[0::2, 0::2] + u[1::2, 0::2] + u[0::2, 1::2] + u[1::2, 1::2]) * 0.25
shrunk_v = (v[0::2, 0::2] + v[1::2, 0::2] + v[0::2, 1::2] + v[1::2, 1::2]) * 0.25


# III. U,V Interleaving
########################
# Size of UV plane is half the number of rows, and same number of columns as Y plane.
uv = np.zeros((rows//2, cols))  # Use // for integer division.

# Interleave shrunkU and shrunkV and build UV plane (each row of UV plane is u,v,u,u,v...)
uv[:, 0::2] = shrunk_u
uv[:, 1::2] = shrunk_v

# Place Y plane at the top, and UV plane at the bottom (number of rows NV12 matrix is rows*1.5)
nv12 = np.vstack((y, uv))

# Round NV12, and cast to uint8.
nv12 = np.round(nv12).astype('uint8')

# Write NV12 array to binary file
nv12.tofile('nv12_output.raw')

# Display NV12 result (display as Grayscale image).
plt.figure()
plt.axis('off')
plt.imshow(nv12, cmap='gray', interpolation='nearest')
plt.show()


# Testing - compare the NV12 result to FFmpeg conversion result:
################################################################################
color_matrix = 'bt709' if do_use_bt709 else 'bt601'

sp.run(['ffmpeg', '-y', '-i', 'rgb_input.png', '-vf', 
        f'scale=flags=fast_bilinear:out_color_matrix={color_matrix}:out_range=tv:dst_format=nv12',
        '-pix_fmt', 'nv12', '-f', 'rawvideo', 'nv12_ffmpeg.raw'])

nv12_ff = np.fromfile('nv12_ffmpeg.raw', np.uint8)
nv12_ff = nv12_ff.reshape(nv12.shape)

abs_diff = np.absolute(nv12.astype(np.int16) - nv12_ff.astype(np.int16)).astype(np.uint8)
max_abs_diff = abs_diff.max()

print(f'max_abs_diff = {max_abs_diff}')

plt.figure()
plt.axis('off')
plt.imshow(abs_diff, cmap='gray', interpolation='nearest')
plt.show()
################################################################################

RGB输入图像示例：

NV12结果（以灰度图像显示）：

---

测试：

为了测试，我们使用命令行工具FFmpeg将相同的输入图像（rgb_input.png）转换为NV12格式，并计算两个转换之间的最大绝对差。

测试假定FFmpeg在执行路径中（在Windows中，我们可以将ffmpeg.exe放置在与Python脚本相同的文件夹中）。

以下shell命令将rgb_input.png转换为BT.709颜色标准的NV12格式：

ffmpeg -y -i rgb_input.png -vf "scale=flags=fast_bilinear:out_color_matrix=bt709:out_range=tv:dst_format=nv12" -pix_fmt nv12 -f rawvideo nv12_ffmpeg.raw

注意：
fast_bilinear插值可在特定输入图像下获得最佳结果-当缩小U和V时应用双线性插值。

以下Python代码将nv12_ffmpeg.raw与nv12_ffmpeg.raw进行比较：

nv12_ff = np.fromfile('nv12_ffmpeg.raw', np.uint8).reshape(nv12.shape)
abs_diff = np.absolute(nv12.astype(np.int16) - nv12_ff.astype(np.int16)).astype(np.uint8)
print(f'max_abs_diff = {abs_diff.max()}')

针对特定的输入图像，最大差异为2或3（几乎相同）。 对于其他输入图像，差异较大（可能是由于错误的FFmpeg参数导致）。

- Rotem