关于YUV420p到YUV444p

关于YUV420p到YUV444p

在项目中遇到一个问题，需要将YUV420p的视频转化为YUV444p的格式，在开始时使用了FFmpeg的命令如下：

ffmpeg -y -s 480x270 -pix_fmt yuv420p -i xx.yuv -f rawvideo -s 480x270 -pix_fmt yuv444p xx-yuv444.yuv

先说结论，上面的命令是正确的，可以将yuv420p正确的转为yuv444p格式。 问题出现在，我怀疑了上面的命令是不是正确。

于是我有了下面的操作，使用OpenCV看一下转化出来的YUV444p是不是对的啊？ 而由于python不能直接读取yuv格式的图像，
所以我需要先以二进制读入文件，参考yuv420p利用OpenCV的读取办法，我写了yuv444p的读取代码。

import numpy as np
import cv2

def read_yuv420p_frame(file_name, dims):
    with open(file_name, "rb") as f:
        h, w = dims
        buffer = f.read(int(h * w * 3 / 2))
        img = np.frombuffer(buffer, dtype=np.uint8).reshape((int(h * 1.5), w))
        return img


def read_yuv444p_frame(file_name, dims):
    with open(file_name, "rb") as f:
        h, w = dims
        buffer = f.read(h * w * 3)
        img = np.frombuffer(buffer, dtype=np.uint8).reshape((h, w, 3))
        return img

yuv420p_file_path = "yuv420p.yuv"
yuv444p_file_path = "yuv444p.yuv"

yuv420p_img = read_yuv420p_frame(yuv420p_file_path, (h, w))
yuv444p_img = read_yuv444p_frame(yuv444p_file_path, (h, w))

cv2.imwrite("yuv420p-to-bgr.png", cv2.cvtColor(yuv420p_img, cv2.COLOR_YUV2BGR_I420))
cv2.imwrite("yuv444p-to-bgr.png", cv2.cvtColor(yuv444p_img, cv2.COLOR_YUV2BGR))

很不幸，虽然yuv420p的代码读取正确，但是yuv444p的代码读取有问题。它的结果变成了这样

为了验证是不是读取的问题，我写了下面的代码，手动将YUV420p转为YUV444p

def convert_yuv420p_to_yuv444p(yuv420p_img, dims):
    # By filling
    img = yuv420p_img.flatten()
    h, w = dims

    U = img[h * w : h * w + (h * w) // 4]
    V = img[h * w + (h * w) // 4 : ]
    U = U.repeat(2).reshape((2, (h * w) // 4)).repeat((2, 2), axis=0).flatten()
    V = V.repeat(2).reshape((2, (h * w) // 4)).repeat((2, 2), axis=0).flatten()
    img = img[: h * w]

    img = np.append(img, U)
    img = np.append(img, V)

    return img.reshape((h, w, 3))

发现还是不对，于是为了进一步验证我手动转化YUV420p到YUV444p是不是有问题，我构造了一个YUV420p的数据（24个Y数据，6个U和6个V），参考wiki
上的图像，最后发现，输出的形状对不上，换句话说reshape写的有问题！

所以正确的代码应该如下，注意最后的输出格式

import numpy as np
import cv2

def read_yuv420p_frame(file_name, dims):
    with open(file_name, "rb") as f:
        h, w = dims
        buffer = f.read(int(h * w * 3 / 2))
        img = np.frombuffer(buffer, dtype=np.uint8).reshape((int(h * 1.5), w))
        return img


def read_yuv444p_frame(file_name, dims):
    with open(file_name, "rb") as f:
        h, w = dims
        buffer = f.read(h * w * 3)
        img = np.frombuffer(buffer, dtype=np.uint8).reshape((3, h, w)).transpose((1, 2, 0))
        return img

def convert_yuv420p_to_yuv444p(yuv420p_img, dims):
    # By filling
    img = yuv420p_img.flatten()
    h, w = dims

    U = img[h * w : h * w + (h * w) // 4]
    V = img[h * w + (h * w) // 4 : ]
    U = U.repeat(2).reshape((2, (h * w) // 4)).repeat((2, 2), axis=0).flatten()
    V = V.repeat(2).reshape((2, (h * w) // 4)).repeat((2, 2), axis=0).flatten()
    img = img[: h * w]

    img = np.append(img, U)
    img = np.append(img, V)

    return img.reshape((3, h, w)).transpose((1, 2, 0))

yuv420p_file_path = "yuv420p.yuv"
yuv444p_file_path = "yuv444p.yuv"

yuv420p_img = read_yuv420p_frame(yuv420p_file_path, (h, w))
yuv444p_img = read_yuv444p_frame(yuv444p_file_path, (h, w))

cv2.imwrite("yuv420p-to-bgr.png", cv2.cvtColor(yuv420p_img, cv2.COLOR_YUV2BGR_I420))
cv2.imwrite("yuv444p-to-bgr.png", cv2.cvtColor(yuv444p_img, cv2.COLOR_YUV2BGR))

此外，在这里记录一下，有一个部分处理的不好，YUV视频提取帧，做的十分暴力，结合python脚本和下面的命令，逐帧提取。

ffmpeg -s widthxheight -i input.yuv -c:v rawvideo -filter:v select="between(n\, 0\, 1)" out.yuv