在Python中将YUV转换为RGB，系数适用于Array，不适用于NumPy

Question

以下代码获取原始 YUV 图像并使用数组操作将其转换为 RGB，并且需要很长时间。

from PIL import Image
from struct import *
import array
 
 
image_name = "frame2.raw" #Change to user input
width = int(3864) #Assumed to be static
height = int(2192) #Assumed to be static

y = array.array('B')  #B unsigned Char, I is unsigned int
u = array.array('B')  #The raw file is unsigned int 8
v = array.array('B')
 
f_y = open(image_name, "rb")
f_uv = open(image_name, "rb")
f_uv.seek(width*height, 1) #Change position of File handle seek(offset, from)
 
image_out = Image.new("RGB", (width, height))
pix = image_out.load()
 
for i in range(0, height//2): #Run height/2 times
    for j in range(0, width//2): #Run width/2 times for each height count
        u.append(ord(f_uv.read(1)));   #append adds to the end of the array; Ord returns the Unicode
        v.append(ord(f_uv.read(1)));   #read byte by byte
 
for i in range(0,height): #Run height times
    for j in range(0, width): #Run width times for each height count
        y.append(ord(f_y.read(1))); #build y byte by byte in unicode from f_y
        Y_val = y[(i*width)+j]
        U_val = u[((i//2)*(width//2))+(j//2)]
        V_val = v[((i//2)*(width//2))+(j//2)]
        B = 1.164 * (Y_val-16) + 2.018 * (U_val - 128)
        G = 1.164 * (Y_val-16) - 0.813 * (V_val - 128) - 0.391 * (U_val - 128)
        R = 1.164 * (Y_val-16) + 1.596*(V_val - 128)
        pix[j, i] = int(R), int(G), int(B) 
 
######################################################
# B = 1.164(Y - 16)                   + 2.018(U - 128)
# G = 1.164(Y - 16) - 0.813(V - 128) - 0.391(U - 128)
# R = 1.164(Y - 16) + 1.596(V - 128)
######################################################
 
image_out.save("out.bmp")

但会产生准确的图像：
基于数组的代码的输出

为了提高效率，我切换到 NumPy，使用相同的颜色转换系数。

from PIL import Image
import numpy as np
 
 
image_name = "frame2.raw" #Change to user input
width = int(3864) #Assumed to be static
height = int(2192) #Assumed to be static
y_end = width*height

yuv = np.fromfile(image_name, dtype='uint8')

y = yuv[0:y_end].reshape(height,width)
u = yuv[y_end::2].reshape(height//2, width//2)
v = yuv[y_end+1::2].reshape(height//2, width//2)

#y = y.astype(y.dtype) -16
#u = u.astype(u.dtype) -128
#v = v.astype(v.dtype) -128
u = u.repeat(2, axis=0).repeat(2, axis=1)
v = v.repeat(2, axis=0).repeat(2, axis=1)

y = y.reshape((y.shape[0], y.shape[1], 1))
u = u.reshape((u.shape[0], u.shape[1], 1))
v = v.reshape((v.shape[0], v.shape[1], 1))

yuv_array = np.concatenate((y, u, v), axis=2)

yuv_array[:, :, 0] = yuv_array[:, :, 0].clip(16, 235).astype(yuv_array.dtype) - 16
yuv_array[:, :, 1:] = yuv_array[:, :, 1:].clip(16, 240).astype(yuv_array.dtype) - 128

convert = np.array([#[1.164,  0.000,  1.793],[1.164, -0.213, -0.533],[1.164,  2.112,  0.000]
                    [1.164,  0.000,  2.018], [1.164, -0.813, -0.391],[1.164,  1.596,  0.000]
                   ])
rgb = np.matmul(yuv_array, convert.T).clip(0,255).astype('uint8')


rgb_image = Image.fromarray(rgb)

rgb_image.save('numpyout.bmp')

但是现在的输出是：
基于 NumPy 的代码的输出

我相信这是我所忽略的简单事情，任何帮助将不胜感激。

链接到原始文件

Answer 1

问题是算术溢出。

看一下下面的代码行：

yuv_array[:, :, 1:] = yuv_array[:, :, 1:].clip(16, 240).astype(yuv_array.dtype) - 128

• yuv_array.dtype是uint8。
  从元素中减去 128 值uint8会导致所有低于 128 的值发生算术溢出。
  示例：
  np.array([100, 100], np.uint8) - 128返回array([228, 228], dtype=uint8)（而不是[-28, -28]）。

解：
转换yuv_array为float32减法前：

yuv_array = yuv_array.astype(np.float32)
yuv_array[:, :, 0] = yuv_array[:, :, 0].clip(16, 235) - 16
yuv_array[:, :, 1:] = yuv_array[:, :, 1:].clip(16, 240) - 128

完整代码：

from PIL import Image
import numpy as np
 
 
image_name = "frame2.raw" #Change to user input
width = int(3864) #Assumed to be static
height = int(2192) #Assumed to be static
y_end = width*height

yuv = np.fromfile(image_name, dtype='uint8')

y = yuv[0:y_end].reshape(height,width)
u = yuv[y_end::2].reshape(height//2, width//2)
v = yuv[y_end+1::2].reshape(height//2, width//2)

u = u.repeat(2, axis=0).repeat(2, axis=1)
v = v.repeat(2, axis=0).repeat(2, axis=1)

y = y.reshape((y.shape[0], y.shape[1], 1))
u = u.reshape((u.shape[0], u.shape[1], 1))
v = v.reshape((v.shape[0], v.shape[1], 1))

yuv_array = np.concatenate((y, u, v), axis=2)

# Overflow: yuv_array.dtype = 'uint8', so subtracting 128 overflows.
#yuv_array[:, :, 0] = yuv_array[:, :, 0].clip(16, 235).astype(yuv_array.dtype) - 16
#yuv_array[:, :, 1:] = yuv_array[:, :, 1:].clip(16, 240).astype(yuv_array.dtype) - 128

# Convert from uint8 to float32 before subtraction
yuv_array = yuv_array.astype(np.float32)
yuv_array[:, :, 0] = yuv_array[:, :, 0].clip(16, 235) - 16
yuv_array[:, :, 1:] = yuv_array[:, :, 1:].clip(16, 240) - 128


convert = np.array([#[1.164,  0.000,  1.793],[1.164, -0.213, -0.533],[1.164,  2.112,  0.000]
                    [1.164,  0.000,  2.018], [1.164, -0.813, -0.391],[1.164,  1.596,  0.000]
                   ])
rgb = np.matmul(yuv_array, convert.T).clip(0,255).astype('uint8')


rgb_image = Image.fromarray(rgb)

rgb_image.save('numpyout.bmp')

测试：

• 使用 FFmpeg 将 JPEG 图像从 RGB 转换为 NV12 格式：（
  ffmpeg -i rgb.jpg -vcodec rawvideo -pix_fmt nv12 -f rawvideo frame2.raw
  我们需要它只是因为我们没有原始frame2.raw文件）。
• 执行更新后的代码。

rgb_image（下采样）：

注意：
您使用的 YUV 到 RGB 转换矩阵适用 BT.601 格式。
您的输入格式更有可能采用 BT.709 格式。