因为在项目中要用到yuv格式的视频图像进行模型推理,但是我们模型通常都是只接收RGB格式的图片,所以在推理前就要先把YUV格式转换为RGB格式。
在代码实现之前,需要去了解什么是YUV格式编码,明白了其编码方式后再看实现会简单的多,如果只是需要代码实现,那可以直接简单略过。
YUV是一种类似RGB的颜色模型,起源于黑白和彩电的过渡时期,是一种颜色编码方法,其中Y代表亮度,UV组合起来可以表示色度。YUV信息只有y的信息就足以显示黑白的图片(所以我们早期的黑白电视就是只用了Y),这样的设计很好地解决了彩色电视机与黑白电视的兼容问题。并且,YUV不像RGB那样要求三个独立的视频信号同时传输,所以用YUV方式传送占用极少的频宽。
YUV码流的存储格式其实与其采样的方式密切相关,主流的采样方式有三种,YUV4:4:4,YUV4:2:2,YUV4:2:0,关于其详细原理,可以参考这篇文章
一文读懂 YUV 的采样与格式
,
YUV与RGB编码
这里就不再过多介绍,本文是实现根据其采样格式来从码流中还原每个像素点的YUV值,然后通过YUV与RGB的转换公式提取出每个像素点的RGB值,然后显示保存下来。
在实现代码之前,首先要明确自己要转换的YUV是什么存储格式,因为不同的存储格式需要不同的提取方式。简单看下下面这两种不同的存储格式:
yuv420P(YU12 和 YV12 )格式如下,先存储Y分量,再存储U分量,最后存储V分量:
yuv420SP(NV12 和 NV21 )格式如下,先存储Y分量,然后交替存储UV分量:
作为对比,再上一张RGB的存储格式:
RGB存储方式:RGB三个分量按照B、G、R的顺序储存。(4:4:4)
所以我们在实现之前需要明确是什么格式存储的,否则提取出来的UV分量可能是错的,导致转换后的图像色彩不对。
第一版实现,使用for循环逐点提取YUV值转换成RGB(非常耗时,不建议使用,为了更清楚实现的过程):
def yuv2rgb(Y, U, V):
bgr_data = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
for h_idx in range(Y_HEIGHT):
for w_idx in range(Y_WIDTH):
y = Y[h_idx, w_idx]
u = U[int(h_idx // 2), int(w_idx // 2)]
v = V[int(h_idx // 2), int(w_idx // 2)]
c = (y - 16) * 298
d = u - 128
e = v - 128
r = (c + 409 * e + 128) // 256
g = (c - 100 * d - 208 * e + 128) // 256
b = (c + 516 * d + 128) // 256
bgr_data[h_idx, w_idx, 2] = 0 if r < 0 else (255 if r > 255 else r)
bgr_data[h_idx, w_idx, 1] = 0 if g < 0 else (255 if g > 255 else g)
bgr_data[h_idx, w_idx, 0] = 0 if b < 0 else (255 if b > 255 else b)
return bgr_data
第二版实现,使用numpy数组运算进行加速(速度非常快,建议用这版):
def np_yuv2rgb(Y,U,V):
bgr_data = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
V = np.repeat(V, 2, 0)
V = np.repeat(V, 2, 1)
U = np.repeat(U, 2, 0)
U = np.repeat(U, 2, 1)
c = (Y-np.array([16])) * 298
d = U - np.array([128])
e = V - np.array([128])
r = (c + 409 * e + 128) // 256
g = (c - 100 * d - 208 * e + 128) // 256
b = (c + 516 * d + 128) // 256
r = np.where(r < 0, 0, r)
r = np.where(r > 255,255,r)
g = np.where(g < 0, 0, g)
g = np.where(g > 255,255,g)
b = np.where(b < 0, 0, b)
b = np.where(b > 255,255,b)
bgr_data[:, :, 2] = r
bgr_data[:, :, 1] = g
bgr_data[:, :, 0] = b
return bgr_data
实验对比:输入同一张(1152*648)的YUV图像。
第一版耗时:5.698601007461548
第二版耗时:0.04670834541320801
速度提升了一百多倍,图片越大,提升效果越明显,最后转换出来的图像如图所示:
全部代码:
import os
import cv2
import numpy as np
IMG_WIDTH = 1152
IMG_HEIGHT = 648
IMG_SIZE = int(IMG_WIDTH * IMG_HEIGHT * 3 / 2)
Y_WIDTH = IMG_WIDTH
Y_HEIGHT = IMG_HEIGHT
Y_SIZE = int(Y_WIDTH * Y_HEIGHT)
U_V_WIDTH = int(IMG_WIDTH / 2)
U_V_HEIGHT = int(IMG_HEIGHT / 2)
U_V_SIZE = int(U_V_WIDTH * U_V_HEIGHT)
def from_I420(yuv_data, frames):
Y = np.zeros((frames, IMG_HEIGHT, IMG_WIDTH), dtype=np.uint8)
U = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
V = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
for frame_idx in range(0, frames):
y_start = frame_idx * IMG_SIZE
u_start = y_start + Y_SIZE
v_start = u_start + U_V_SIZE
v_end = v_start + U_V_SIZE
Y[frame_idx, :, :] = yuv_data[y_start : u_start].reshape((Y_HEIGHT, Y_WIDTH))
U[frame_idx, :, :] = yuv_data[u_start : v_start].reshape((U_V_HEIGHT, U_V_WIDTH))
V[frame_idx, :, :] = yuv_data[v_start : v_end].reshape((U_V_HEIGHT, U_V_WIDTH))
return Y, U, V
def from_YV12(yuv_data, frames):
Y = np.zeros((frames, IMG_HEIGHT, IMG_WIDTH), dtype=np.uint8)
U = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
V = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
for frame_idx in range(0, frames):
y_start = frame_idx * IMG_SIZE
v_start = y_start + Y_SIZE
u_start = v_start + U_V_SIZE
u_end = u_start + U_V_SIZE
Y[frame_idx, :, :] = yuv_data[y_start : v_start].reshape((Y_HEIGHT, Y_WIDTH))
V[frame_idx, :, :] = yuv_data[v_start : u_start].reshape((U_V_HEIGHT, U_V_WIDTH))
U[frame_idx, :, :] = yuv_data[u_start : u_end].reshape((U_V_HEIGHT, U_V_WIDTH))
return Y, U, V
def from_NV12(yuv_data, frames):
Y = np.zeros((frames, IMG_HEIGHT, IMG_WIDTH), dtype=np.uint8)
U = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
V = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
for frame_idx in range(0, frames):
y_start = frame_idx * IMG_SIZE
u_v_start = y_start + Y_SIZE
u_v_end = u_v_start + (U_V_SIZE * 2)
Y[frame_idx, :, :] = yuv_data[y_start : u_v_start].reshape((Y_HEIGHT, Y_WIDTH))
U_V = yuv_data[u_v_start : u_v_end].reshape((U_V_SIZE, 2))
U[frame_idx, :, :] = U_V[:, 0].reshape((U_V_HEIGHT, U_V_WIDTH))
V[frame_idx, :, :] = U_V[:, 1].reshape((U_V_HEIGHT, U_V_WIDTH))
return Y, U, V
def from_NV21(yuv_data, frames):
Y = np.zeros((frames, IMG_HEIGHT, IMG_WIDTH), dtype=np.uint8)
U = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
V = np.zeros((frames, U_V_HEIGHT, U_V_WIDTH), dtype=np.uint8)
for frame_idx in range(0, frames):
y_start = frame_idx * IMG_SIZE
u_v_start = y_start + Y_SIZE
u_v_end = u_v_start + (U_V_SIZE * 2)
Y[frame_idx, :, :] = yuv_data[y_start : u_v_start].reshape((Y_HEIGHT, Y_WIDTH))
U_V = yuv_data[u_v_start : u_v_end].reshape((U_V_SIZE, 2))
V[frame_idx, :, :] = U_V[:, 0].reshape((U_V_HEIGHT, U_V_WIDTH))
U[frame_idx, :, :] = U_V[:, 1].reshape((U_V_HEIGHT, U_V_WIDTH))
return Y, U, V
def np_yuv2rgb(Y,U,V):
bgr_data = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
V = np.repeat(V, 2, 0)
V = np.repeat(V, 2, 1)
U = np.repeat(U, 2, 0)
U = np.repeat(U, 2, 1)
c = (Y-np.array([16])) * 298
d = U - np.array([128])
e = V - np.array([128])
r = (c + 409 * e + 128) // 256
g = (c - 100 * d - 208 * e + 128) // 256
b = (c + 516 * d + 128) // 256
r = np.where(r < 0, 0, r)
r = np.where(r > 255,255,r)
g = np.where(g < 0, 0, g)
g = np.where(g > 255,255,g)
b = np.where(b < 0, 0, b)
b = np.where(b > 255,255,b)
bgr_data[:, :, 2] = r
bgr_data[:, :, 1] = g
bgr_data[:, :, 0] = b
return bgr_data
def yuv2rgb(Y, U, V):
bgr_data = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
for h_idx in range(Y_HEIGHT):
for w_idx in range(Y_WIDTH):
y = Y[h_idx, w_idx]
u = U[int(h_idx // 2), int(w_idx // 2)]
v = V[int(h_idx // 2), int(w_idx // 2)]
c = (y - 16) * 298
d = u - 128
e = v - 128
r = (c + 409 * e + 128) // 256
g = (c - 100 * d - 208 * e + 128) // 256
b = (c + 516 * d + 128) // 256
bgr_data[h_idx, w_idx, 2] = 0 if r < 0 else (255 if r > 255 else r)
bgr_data[h_idx, w_idx, 1] = 0 if g < 0 else (255 if g > 255 else g)
bgr_data[h_idx, w_idx, 0] = 0 if b < 0 else (255 if b > 255 else b)
return bgr_data
if __name__ == '__main__':
import time
yuv = "request/YUV/2021-05-06/test.yuv"
frames = int(os.path.getsize(yuv) / IMG_SIZE)
with open(yuv, "rb") as yuv_f:
time1 = time.time()
yuv_bytes = yuv_f.read()
yuv_data = np.frombuffer(yuv_bytes, np.uint8)
Y, U, V = from_NV21(yuv_data, frames)
rgb_data = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
for frame_idx in range(frames):
bgr_data = np_yuv2rgb(Y[frame_idx, :, :], U[frame_idx, :, :], V[frame_idx, :, :])
time2 = time.time()
print(time2 - time1)
if bgr_data is not None:
cv2.imwrite("frame_{}.jpg".format(frame_idx), bgr_data)
frame_idx +=1
代码和测试用的YUV文件已上传到github,使用时请参考文档。
序言因为在项目中要用到yuv格式的视频图像进行模型推理,但是我们模型通常都是只接收RGB格式的图片,所以在推理前就要先把YUV格式转换为RGB格式,网上搜了看到很多实现,搜来搜去你会发现还是那几个源码,copy下来运行要么报错,要么转出来的色彩不对,全都是没经过实践的,抄袭得特别严重,一份代码写了十几篇文章,关键是还不一定能用,来回调试严重浪费我们宝贵的时间,再次对这种行为表示愤慨。如果你一开始就搜到了我这篇文章,那么恭喜你,不用再去踩我曾经踩过的坑。一、什么是YUV在代码实现之前,需要去了解什么是
一、将base64的图片数据转为RGB (速度很慢)
因为使用的Python内置的for循环遍历图片数据,所以速度会比较慢,转换一张200k左右的图片需要2s到3s左右。
#encoding:utf-8
import numpy as np
import cv2
video_cutoff = [0.0625, 0.5, 0.5]
full_cutoff = [0.0, 0.5, 0.5]...
应用:模拟领域
Y'= 0.299*R' + 0.587*G' + 0.114*B'
U'= -0.147*R' - 0.289*G' + 0.436*B' = 0.492*(B'- Y')
V'= 0.615*R' - 0.
支持yuv\rgb的各种处理,比如rgb与yuv的相互转换,视频格式转换,从一段长视频中截取出某几帧等,此工具包里包含多个工具
YUVviewerPlus.exe
对YUVviewer进行了修改,增加一下内容:
1、增加支持的格式:yuv4:4:4, yuv4:2:2, gbmp
2、增加zoom的范围
BMP2GBMP.exe
将一副副的BMP图片合并成一个没有BMP头信息的文件
note:仅支持24位bmp图片
DYUV2SEQ.exe
实现分离的yuv文件转换成YUV序列 4:2:0
ShowDIB.exe
BMP图片显示程序,多文档框架
YUV2BMP.exe
实现了YUV转换成24位的BMP图片,实现了批量转换
BMP2SEQ.exe
将一系列24位或8位的BMP图片转换成4:2:0的YUV序列。
DYUV2BMP.exe
将分离的Y,U,V转换成24位的BMP图片,实现了批量转换
GBMP2SEQ.exe
实现包含一组rgb24图片(不含文件头)的单一文件到yuv4:2:0序列文件的转换
SEQ2BMP.exe
实现了SEQ2BMP的程序
输出BMP文件为24位真彩
SeqCut.exe
实现对YUV4:2:0文件的剪切操作
即从序列文件中取出一段序列
SeqSnr.exe
实现了两序列对应帧之间Y分量的SNR求取,并给出平均值
YUV2SEQ.exe
将单帧的YUV文件转换位YUV序列 4:2:0
YUV2SEQ2.exe
将单帧的YUV文件转换位YUV序列 4:2:0
可以选择目标图像的位置和大小
y, u, v = yuv[0], yuv[1] - 128, yuv[2] - 128
r = y + 1.13983 * v
g = y - 0.39465 * u - 0.58060 * v
b = y + 2.03211 * u
return int(r), int(g), int(b)
其中,yuv是一个长度为3的元组或列表,分别表示图像中的Y、U、V分量。返回值为一个长度为3的元组,分别表示RGB图像中的R、G、B分量。请注意,该函数仅适用于YUV格式为YUV444的图像。如果您的图像格式为YUV422或YUV420,请使用适当的转换算法。