python YUV420视频转图像帧

YUV420p视频转成多帧图像

import os
import cv2
import numpy as np
from typing import List

def YUVvideo2IMGs(file, savepath, height, width ):
    """
    Convert the YUV video to RGB images and save the images at the target folder
    Args:
        file: input yuv file name
        savepath: save path of each RGB images
        height: height of images
        width: width of images

    Returns:

    """
    img_size = int(width * height * 3 / 2)

    frames = int(os.path.getsize(file) / img_size)
    with open(file, "rb") as f:
        for frame_idx in range(frames):
            yuv = np.zeros(shape=img_size, dtype='uint8', order='C')
            for j in range(img_size):
                yuv[j] = ord(f.read(1))
            img = yuv.reshape((height * 3 // 2, width)).astype('uint8')
            bgr_img = cv2.cvtColor(img, cv2.COLOR_YUV2BGR_I420)
            if bgr_img is not None:
                cv2.imwrite(os.path.join(savepath, "f%03d.png" % (frame_idx+1)), bgr_img)

多帧图像转成YUV420p

import os
import cv2
import numpy as np
from typing import List

IMG_EXTENSIONS = (
    ".jpg",
    ".jpeg",
    ".png",
    ".ppm",
    ".bmp",
    ".pgm",
    ".tif",
    ".tiff",
    ".webp",
)

def collect_images(rootpath: str) -> List[str]:
    return [
        os.path.join(rootpath, f)
        for f in os.listdir(rootpath)
        if os.path.splitext(f)[-1].lower() in IMG_EXTENSIONS
    ]

def IMG2YUVvideo(Imgdir, YUVfile):
    """

    :param Imgdir: input image folder of the
    :param YUVfile:
    :return:
    """
    Imagepaths = sorted(collect_images(Imgdir))
    img = cv2.imread(Imagepaths[0])
    (height, width, _) = img.shape
    with open(YUVfile, "wb+") as f:
        for i in Imagepaths:
            img = cv2.imread(i)
            YUV = cv2.cvtColor(img, cv2.COLOR_BGR2YUV_I420)
            YUV = YUV.reshape(height * width * 3 // 2)
            f.write(np.array(YUV, dtype=np.uint8).tobytes())