YUV视频序列转RGB图像

# %% md

## yuv视频转为rgb图像
""""
将一个文件夹中包含不同场景的文件夹下的YUV转为图像，对应保存到另外一个文件夹，
相同的该文件夹下有着代表不同场景的文件夹，这些文件夹下的图像就是由YUV转来的
"""

# %%

import os
import cv2
import numpy as np


def yuv2bgr(video_dir, height, width, startfrm):
    """
    :param filename: 待处理 YUV 视频的名字
    :param height: YUV 视频中图像的高
    :param width: YUV 视频中图像的宽
    :param startfrm: 起始帧
    :return: None
    """
    """A dataset for loading filelist as.
        like::
            data/case1/1.jpg
            data/case1/2.jpg
            data/case1/N.jpg
            data/case2/1.jpg
            data/case2/2.jpg
            data/case2/N.jpg
        """
    caselist = os.listdir(video_dir)
    # 转换test文件下的yuv时关闭下面这句话
    # caselist.sort(key=lambda x: int(x.split('_')[0]))

    for case in caselist:
        filepath = os.path.join(video_dir, case)
        print(filepath)
        filenames = os.listdir(filepath)
        filenames.sort(key=lambda x: int((x.split('_', 5)[5]).split('.')[0]))
        print(filenames)
        for filename in filenames:
            fp = open(os.path.join(filepath, filename), 'rb')

            # YUV420是RGB24内存的一半
            framesize = height * width * 3 // 2  # 一帧图像所含的像素个数
            h_h = height // 2
            h_w = width // 2

            fp.seek(0, 2)  # 设置文件指针到文件流的尾部
            ps = fp.tell()  # 当前文件指针位置
            numfrm = ps // framesize  # 计算输出帧数
            fp.seek(framesize * startfrm, 0)

            for i in range(numfrm - startfrm):
                Yt = np.zeros(shape=(height, width), dtype='uint8', order='C')
                Ut = np.zeros(shape=(h_h, h_w), dtype='uint8', order='C')
                Vt = np.zeros(shape=(h_h, h_w), dtype='uint8', order='C')

                for m in range(height):
                    for n in range(width):
                        # ord以一个字符（长度为1的字符串）作为参数，返回对应的 ASCII 数值，或者 Unicode 数值
                        Yt[m, n] = ord(fp.read(1))
                for m in range(h_h):
                    for n in range(h_w):
                        Ut[m, n] = ord(fp.read(1))
                for m in range(h_h):
                    for n in range(h_w):
                        Vt[m, n] = ord(fp.read(1))

                img = np.concatenate((Yt.reshape(-1), Ut.reshape(-1), Vt.reshape(-1)))

                img = img.reshape((height * 3 // 2, width)).astype('uint8')  # YUV 的存储格式为：NV12（YYYY UV）

                # 由于 opencv 不能直接读取 YUV 格式的文件, 所以要转换一下格式
                bgr_img = cv2.cvtColor(img, cv2.COLOR_YUV2BGR_I420)  # 注意 YUV 的存储格式

                save_dir = "/media/ps/2tb/yjz/data/our_data/test/Temporary/rgb/{}".format(case)
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                save_path = os.path.join(save_dir, '{}.png'.format(filename.split('.')[0]))

                cv2.imwrite(save_path, bgr_img)
            fp.close()
    print("finished")
    return None


if __name__ == '__main__':
    _ = yuv2bgr(video_dir="yuv路径", height=1080, width=1920, startfrm=0)