Python OpenCV将平面YUV 4:2:0图像转换为RGB -- YUV数组格式

5
我正在尝试使用OpenCV,通过Python将平面YUV 4:2:0图像转换为RGB,并且在理解如何格式化数组以传递给函数方面遇到了困难。我有所有三个通道作为单独的数组,并尝试将它们合并以供cv2.cvtColor使用。我正在使用cv2.cvtColor(yuv_array, cv2.COLOR_YUV420p2RGB)。我了解到yuv_array应该比原始图像高1.5倍(这就是使用cv2.COLOR_RGB2YUV_YV12从得到的yuv数组的样子),并且我应该将UV分量放入yuv_array的底部一半,将Y通道放入数组的顶部。

我似乎无法弄清楚U和V通道应该如何在该数组的底部格式化。我已经尝试过交错它们并将它们都放在那里。对于两种方法,我都尝试先放U再放V,也尝试了反过来。所有方法都会导致结果图像中出现伪影。这是我的代码和一个示例图像:

import os
import errno
import numpy as np
import cv2

fifo_names = ["/tmp/fifos/y_fifo", "/tmp/fifos/u_fifo", "/tmp/fifos/v_fifo"]

#teardown; delete fifos
import signal, sys
def cleanup_exit(signal, frame):
    print ("cleaning up!")
    for fifo in fifo_names:
        os.remove(fifo)
    sys.exit(0)
signal.signal(signal.SIGINT, cleanup_exit)
signal.signal(signal.SIGTERM, cleanup_exit)

#make fifos
for fifo in fifo_names:
    try:
        os.mkfifo(fifo);
    except OSError as oe:
        if oe.errno == errno.EEXIST:
            os.remove(fifo)
            os.mkfifo(fifo)
        else:
            raise()

#make individual np arrays to store Y,U, and V channels
#we know the image size beforehand -- 640x360 pixels
yuv_data = []
frame_size = []
fullsize = (360, 640)
halfsize = (180, 320)
for i in range(len(fifo_names)):
    if (i == 0):
        size = fullsize
    else:
        size = halfsize
    yuv_data.append(np.empty(size, dtype=np.uint8));
    frame_size.append(size)

#make array that holds all yuv data for display with cv2
all_yuv_data = np.empty((fullsize[0] + halfsize[0], fullsize[1]), dtype=np.uint8) 

#continuously read yuv images from fifos
print("waiting for fifo to be written to...")
while True:
    for i in range(len(fifo_names)):
        fifo = fifo_names[i]
        with open(fifo, 'rb') as f:
            print("FIFO %s opened" % (fifo))
            all_data = b''
            while True:
                data = f.read()
                print("read from %s, len: %d" % (fifo,len(data)))
                if len(data) == 0: #then the fifo has been closed
                    break
                else:
                    all_data += data
            yuv_data[i] = np.frombuffer(all_data, dtype=np.uint8).reshape(frame_size[i])

    #stick all yuv data in one buffer, interleaving columns
    all_yuv_data[0:fullsize[0],0:fullsize[1]] = yuv_data[0]
    all_yuv_data[fullsize[0]:,0:fullsize[1]:2] = yuv_data[1]
    all_yuv_data[fullsize[0]:,1:fullsize[1]:2] = yuv_data[2]

    #show each yuv channel individually
    cv2.imshow('y', yuv_data[0])
    cv2.imshow('u', yuv_data[1])
    cv2.imshow('v', yuv_data[2])

    #convert yuv to rgb and display it
    rgb = cv2.cvtColor(all_yuv_data, cv2.COLOR_YUV420p2RGB);
    cv2.imshow('rgb', rgb)
    cv2.waitKey(1)

上面的代码试图以列为单位交错U和V信息。
我还尝试使用以下方法将U和V通道信息放入all_yuv_data数组中:
    #try back-to-back
    all_yuv_data[0:fullsize[0],0:fullsize[1]] = yuv_data[0]
    all_yuv_data[fullsize[0]:,0:halfsize[1]] = yuv_data[1]
    all_yuv_data[fullsize[0]:,halfsize[1]:] = yuv_data[2]

这是一幅由libav从另一个程序中获取的视频帧图像,其格式为AV_PIX_FMT_YUV420P被描述为“平面YUV 4:2:0,12bpp,(每2x2个Y样本对应1个Cr和Cb样本)”。
以下是灰度显示的示例图像的yuv通道: Y通道:

y channel

U通道:

u channel

V频道:

v channel

以下是使用上述交错方法得到的对应RGB转换结果(使用“背靠背”方法时也会出现类似伪影):

带有伪影的RGB图像:

rgb image with artifacts

我应该如何在all_yuv_data中放置u和v通道信息? Mark Setchell编辑后的内容: 我认为期望的结果是:

enter image description here


@MarkSetchell 感谢您的反馈,我已经用实际图片替换了屏幕截图。 - founta
我已经添加了我认为是期望结果的内容 - 请告诉我如果我错了,我会撤销我的编辑。谢谢。 - Mark Setchell
@MarkSetchell 是的,那正是我期望的结果,谢谢您添加它。您是如何获得它的? - founta
ImageMagick命令为{ convert y.png -depth 8 gray:- ; convert u.png -depth 8 gray:- ; convert v.png -depth 8 gray:- ; } | convert -depth 8 -size 640x360 yuv:- result.jpg - Mark Setchell
今晚太晚了,所以我明天会检查一下HansHirse、Nathancy、Rotem或Fred是否已经完成了它。 - Mark Setchell
显示剩余3条评论
2个回答

10
如果YUV标准与OpenCV的COLOR_YUV2BGR_I420转换公式匹配,您可以一次性读取帧并将其重塑为高度*1.5行应用转换。
以下代码示例:
构建YUV420格式的输入,并将其写入内存流中(而不是fifo)。
从流中读取帧并使用COLOR_YUV2BGR_I420将其转换为BGR。颜色不正确...
通过读取Y、U和V,调整U和V的大小,并使用COLOR_YCrCb2BGR转换重复此过程。注意:OpenCV在BGR颜色格式下工作(而不是RGB)。
这是代码:
import cv2
import numpy as np
import io

# Building the input:
###############################################################################
img = cv2.imread('GrandKingdom.jpg')

#yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
#y, u, v = cv2.split(yuv)

# Convert BGR to YCrCb (YCrCb apply YCrCb JPEG (or YCC), "full range", 
# where Y range is [0, 255], and U, V range is [0, 255] (this is the default JPEG format color space format).
yvu = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
y, v, u = cv2.split(yvu)

# Downsample U and V (apply 420 format).
u = cv2.resize(u, (u.shape[1]//2, u.shape[0]//2))
v = cv2.resize(v, (v.shape[1]//2, v.shape[0]//2))

# Open In-memory bytes streams (instead of using fifo)
f = io.BytesIO()

# Write Y, U and V to the "streams".
f.write(y.tobytes())
f.write(u.tobytes())
f.write(v.tobytes())

f.seek(0)
###############################################################################

# Read YUV420 (I420 planar format) and convert to BGR
###############################################################################
data = f.read(y.size*3//2)  # Read one frame (number of bytes is width*height*1.5).

# Reshape data to numpy array with height*1.5 rows
yuv_data = np.frombuffer(data, np.uint8).reshape(y.shape[0]*3//2, y.shape[1])

# Convert YUV to BGR
bgr = cv2.cvtColor(yuv_data, cv2.COLOR_YUV2BGR_I420);


# How to How should I be placing the u and v channel information in all_yuv_data?
# -------------------------------------------------------------------------------
# Example: place the channels one after the other (for a single frame)
f.seek(0)
y0 = f.read(y.size)
u0 = f.read(y.size//4)
v0 = f.read(y.size//4)
yuv_data = y0 + u0 + v0
yuv_data = np.frombuffer(yuv_data, np.uint8).reshape(y.shape[0]*3//2, y.shape[1])
bgr = cv2.cvtColor(yuv_data, cv2.COLOR_YUV2BGR_I420);
###############################################################################

# Display result:
cv2.imshow("bgr incorrect colors", bgr)


###############################################################################
f.seek(0)
y = np.frombuffer(f.read(y.size), dtype=np.uint8).reshape((y.shape[0], y.shape[1]))  # Read Y color channel and reshape to height x width numpy array
u = np.frombuffer(f.read(y.size//4), dtype=np.uint8).reshape((y.shape[0]//2, y.shape[1]//2))  # Read U color channel and reshape to height x width numpy array
v = np.frombuffer(f.read(y.size//4), dtype=np.uint8).reshape((y.shape[0]//2, y.shape[1]//2))  # Read V color channel and reshape to height x width numpy array

# Resize u and v color channels to be the same size as y
u = cv2.resize(u, (y.shape[1], y.shape[0]))
v = cv2.resize(v, (y.shape[1], y.shape[0]))
yvu = cv2.merge((y, v, u)) # Stack planes to 3D matrix (use Y,V,U ordering)

bgr = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2BGR)
###############################################################################


# Display result:
cv2.imshow("bgr", bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()

结果:
在此输入图片描述


谢谢你说出来。我的直觉告诉我,OP在这里做了太多的工作。 - Mad Physicist
我更新了我的帖子,COLOR_YCrCb2BGR 的颜色格式是YCrCb JPEG(或YCC) - Rotem

3
在这个函数调用中,存储在 yuv_array 底部的 u 和 v 通道信息应按以下格式进行格式化: cv2.cvtColor(yuv_array, cv2.COLOR_YUV420p2RGB) 它们应该按照以下方式填充:
1. 添加到 yuv_array 底部的额外行的上半部分填充了 u 信息。这些行是交错的;第一行 u 放置在左侧插槽中的 y 通道信息正下方,第二行 u 放置在同一行中的右侧插槽中,以此类推。 2. v 通道数据也是如此,但对于添加到 yuv_array 的底部半部分的额外行。
这是 MarkSetchnell 发布的期望图像所使用的串联代码:
    #place y channel into buffer
    all_yuv_data[0:fullsize[0],0:fullsize[1]] = yuv_data[0]

    #formatted as interleaved u rows on top, (half on left, half on right)
    #and interleaved v rows on bottom
    all_yuv_data[fullsize[0]:fullsize[0]+halfsize[0]//2, :] = yuv_data[1].reshape(-1, fullsize[1])
    all_yuv_data[fullsize[0]+halfsize[0]//2:,:] = yuv_data[2].reshape(-1, fullsize[1])

    #convert to rgb
    rgb = cv2.cvtColor(all_yuv_data, cv2.COLOR_YUV420p2RGB);

这是一张 all_yuv_data 的灰度图像,为了更好地表达,如下所示: 以 y 通道在上、u 在中间、v 在下的 all_yuv_data 灰度图像 调用 cv2.cvtColor(all_yuv_data, cv2.COLOR_YUV420p2RGB) 后得到的结果是: 正确着色的 Grand Kingdom 启动画面美丽的图像

1
你可以通过执行 all_yuv_data[fullsize[0]:fullsize[0]+halfsize[0], :] = yuv_data[1].reshape(-1, fullsize[1]) 以及类似的操作来使代码更易读。对于底层线性缓冲区的顺序,这才是你真正关心的。同时也要对最后一个通道进行类似的操作。 - Mad Physicist
1
正如您所看到的,颜色并不准确。有太多的颜色格式... - Rotem
@Rotem 哦,嘿,你说得对。我甚至没有注意到我的颜色有问题。谢谢你指出来。 - founta
@MadPhysicist 好的,谢谢。我假设你的意思是 all_yuv_data[fullsize[0]:fullsize[0]+halfsize[0]//2, :] = yuv_data[1].reshape(-1, fullsize[1]) - founta
是的,你说得对。我心里想的是实际上使用一个视图来查看缓冲区,例如 small = all_yuv_data[fullsize[0]:, :].reshape(2, -1, halfsize[1]); small[0] = yuv_data[1]; small[1] = yuv_data[2]。这就是当你不干扰连续性时,numpy视图的魔力所在。如果你担心重塑会在途中制作副本,你可以通过使用适当的偏移量和步幅在同一缓冲区上创建一个新数组来明确地创建视图。 - Mad Physicist

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接