Как преобразовать sRGB в формат NV12 с помощью NumPy?

Преобразование sRGB в формат NV12 с помощью NumPy

Цель поста — продемонстрировать процесс преобразования.
В приведенной ниже реализации Python используется NumPy и намеренно не используется OpenCV.

Этапы преобразования RGB в NV12:

• Преобразование цветового пространства — преобразование из цветового пространства sRGB в цветовое пространство YUV:
  Используйте формулу преобразования sRGB в YCbCr.
  Умножьте каждую тройку RGB на матрицу преобразования 3x3 и добавьте вектор из 3 смещений.
  В сообщении показаны оба BT Преобразования .709 и BT.601 (разница только в матрице коэффициентов).
• Понижение дискретизации цветности – сжатие каналов U, V в 2 раза по каждой оси (преобразование из YUV444 в YUV420).
  Реализация изменяет размер U, V в 0,5 раза по каждой оси с использованием билинейной интерполяции.
  Примечание. : билинейная интерполяция не является оптимальным методом понижения частоты дискретизации, но обычно она достаточно хороша.
  Вместо использования cv2.resize код использует среднее значение для каждых 2x2 пикселей (результат эквивалентен билинейной интерполяции).
  Примечание: реализация терпит неудачу в случае, если входное разрешение не даже в обоих измерениях.
• Чередование элементов цветности - расположите элементы U, V как U, V, U, V...
  Реализовано путем манипулирования индексацией массива.

Вот пример кода Python для преобразования RGB в стандарт NV12:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

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 = RGB[:, :, 0], RGB[:, :, 1], RGB[:, :, 2]  # 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.18258588*R + 0.61423059*G + 0.06200706*B + 16.0
    U = -0.10064373*R - 0.33857195*G + 0.43921569*B + 128.0
    V =  0.43921569*R - 0.39894216*G - 0.04027352*B + 128.0
else:
    # Convert sRGB to YUV, BT.601 standard.
    # Conversion formula used: 8 bit sRGB to "limited range" 8 bit YUV (BT.601).
    Y =  0.25678824*R + 0.50412941*G + 0.09790588*B + 16.0
    U = -0.14822290*R - 0.29099279*G + 0.43921569*B + 128.0
    V =  0.43921569*R - 0.36778831*G - 0.07142737*B + 128.0


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

# Each element of shrunkU is the mean of 2x2 elements of U
# Result is equvalent to resize by a factor of 0.5 with bi-linear interpolation.
shrunkU = (U[0: :2, 0::2] + U[1: :2, 0: :2] + U[0: :2, 1: :2] + U[1: :2, 1: :2]) * 0.25
shrunkV = (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 palne (each row of UV plane is u,v,u,u,v...)
UV[:, 0 : :2] = shrunkU
UV[:, 1 : :2] = shrunkV

# 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 (use floor(x+0.5) instead of round to avoid "bankers rounding").
NV12 = np.floor(NV12 + 0.5).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()

Пример входного изображения RGB:

Результат NV12 (отображается как изображение в градациях серого):