在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的代码输出
我相信我可能遗漏了一些简单的东西，如果有帮助，将不胜感激。 原始文件链接

- Gregg Diffendale

---

请分享原始输入文件，以供大家进行实验 - 可以使用Dropbox或Google Drive。还需要提供它的宽度和高度。 - Mark Setchell

宽度为3864，高度为2192。链接 - Gregg Diffendale

Answer 1

问题出在算术溢出上。

看一下以下代码行：

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

• yuv_array.dtype 是uint8类型。
  从128以下的所有值中减去128将导致算术溢出。
  例如:
  np.array([100, 100], np.uint8) - 128的结果为array([228, 228], dtype=uint8)（而不是[-28, -28]）。

---

Solution:
在减法运算之前将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将RGB格式的JPEG图像转换为NV12格式：
  ffmpeg -i rgb.jpg -vcodec rawvideo -pix_fmt nv12 -f rawvideo frame2.raw
  （我们之所以需要这样做，是因为我们没有原始的frame2.raw文件）。
• 执行更新后的代码。

---

rgb_image（降采样）：

---

注意：
您正在使用的YUV到RGB转换矩阵应用BT.601格式。
更有可能的是，您的输入格式应用了BT.709格式。

- Rotem

---

干得好！你有我的支持。 - Mark Setchell

太棒了，谢谢。我添加了一个链接到原始的.raw文件，这样每个人都可以得到一些解决方案。我会检查格式，再次感谢。 - Gregg Diffendale