用python读取YUV文件 转RGB 8bit/10bit通用

注：本文所指的YUV均为YUV420中的I420格式（最常见的一种），其他格式不能用以下的代码。

位深为8bit时，每个像素占用1字节，对应文件指针的fp.read(1)；

位深为10bit时，每个像素占用2字节，对应文件指针的fp.read(2)；

然后使用 int.from_bytes() 方法将二进制转换为int型数字。

以下程序可以读8bit或10bit位深的YUV，需要指定从第几帧开始读、一共读多少帧。

它返回三个数组，其shape分别为：Y [frame,W,H]   U [frame,W/2,H/2]   V [frame,W/2,H/2]

当只读1帧时它返回：Y [W,H]   U [W/2,H/2]   V [W/2,H/2]

# -*- coding: utf-8 -*-

import math
from functools import partial
import numpy as np
import matplotlib.pyplot as plt


def readyuv420(filename, bitdepth, W, H, startframe, totalframe, show=False):
    # 从第startframe（含）开始读（0-based），共读totalframe帧

    uv_H = H // 2
    uv_W = W // 2

    if bitdepth == 8:
        Y = np.zeros((totalframe, H, W), np.uint8)
        U = np.zeros((totalframe, uv_H, uv_W), np.uint8)
        V = np.zeros((totalframe, uv_H, uv_W), np.uint8)
    elif bitdepth == 10:
        Y = np.zeros((totalframe, H, W), np.uint16)
        U = np.zeros((totalframe, uv_H, uv_W), np.uint16)
        V = np.zeros((totalframe, uv_H, uv_W), np.uint16)

    plt.ion()

    bytes2num = partial(int.from_bytes, byteorder='little', signed=False)

    bytesPerPixel = math.ceil(bitdepth / 8)
    seekPixels = startframe * H * W * 3 // 2
    fp = open(filename, 'rb')
    fp.seek(bytesPerPixel * seekPixels)

    for i in range(totalframe):

        for m in range(H):
            for n in range(W):
                if bitdepth == 8:
                    pel = bytes2num(fp.read(1))
                    Y[i, m, n] = np.uint8(pel)
                elif bitdepth == 10:
                    pel = bytes2num(fp.read(2))
                    Y[i, m, n] = np.uint16(pel)

        for m in range(uv_H):
            for n in range(uv_W):
                if bitdepth == 8:
                    pel = bytes2num(fp.read(1))
                    U[i, m, n] = np.uint8(pel)
                elif bitdepth == 10:
                    pel = bytes2num(fp.read(2))
                    U[i, m, n] = np.uint16(pel)

        for m in range(uv_H):
            for n in range(uv_W):
                if bitdepth == 8:
                    pel = bytes2num(fp.read(1))
                    V[i, m, n] = np.uint8(pel)
                elif bitdepth == 10:
                    pel = bytes2num(fp.read(2))
                    V[i, m, n] = np.uint16(pel)

        if show:
            print(i)
            plt.subplot(131)
            plt.imshow(Y[i, :, :], cmap='gray')
            plt.subplot(132)
            plt.imshow(U[i, :, :], cmap='gray')
            plt.subplot(133)
            plt.imshow(V[i, :, :], cmap='gray')
            plt.show()
            plt.pause(1)
            #plt.pause(0.001)

    if totalframe==1:
        return Y[0], U[0], V[0]
    else:
        return Y,U,V


if __name__ == '__main__':
    #y, u, v = readyuv420(r'F:\_commondata\video\176x144 qcif\football_qcif.yuv', 8, 176, 144, 1, 5, True)
    y, u, v = readyuv420(r'F:\_commondata\video\1920x1080 B\RitualDance_1920x1080_60fps_10bit_420.yuv', 10, 1920, 1080, 0, 5, True)
    print(y.shape,u.shape,v.shape)

以下程序将YUV转为RGB（只能读8bit位深的YUV），返回1个数组，其shape为： [frame,W,H,3]

# -*- coding: utf-8 -*-
import cv2
import numpy as np
import matplotlib.pyplot as plt


def readyuv(filename, width, height, framenum, show=False):
    arr = np.zeros((framenum,height,width,3),np.uint8)
    
    plt.ion()
    with open(filename, 'rb') as fp:
        for i in range(framenum):
            print(i)
            oneframe_I420 = np.zeros((height*3//2,width),np.uint8)
            for j in range(height*3//2):
                for k in range(width):
                    oneframe_I420[j,k] = int.from_bytes(fp.read(1), byteorder='little', signed=False)
            oneframe_RGB = cv2.cvtColor(oneframe_I420,cv2.COLOR_YUV2RGB_I420)
            arr[i] = oneframe_RGB
            if show:
                plt.imshow(oneframe_RGB)
                plt.show()
                plt.pause(0.001)
    return arr


if __name__ == '__main__':
    video = readyuv(r'F:\_commondata\video\176x144 qcif\football_qcif.yuv', 176, 144, 130, True)