python提取YUV文件三个分量，加写回操作

# 提取单帧
import numpy as np
from PIL import Image
raw_yuv_path = r'C:\Users\weiwei\Desktop\0.yuv'
raw_write_path = r'C:\Users\weiwei\Desktop\1.yuv'
h, w, nfs = 240, 416, 50
def import_yuv(seq_path, h, w, tot_frm, yuv_type='420p', start_frm=0, only_y=True):
    if yuv_type == '420p':
        hh, ww = h // 2, w // 2
    elif yuv_type == '444p':
        hh, ww = h, w
    else:
        raise Exception('yuv_type not supported.')

    y_size, u_size, v_size = h * w, hh * ww, hh * ww
    blk_size = y_size + u_size + v_size

    # init
    y_seq = np.zeros((tot_frm, h, w), dtype=np.uint8)
    if not only_y:
        u_seq = np.zeros((tot_frm, hh, ww), dtype=np.uint8)
        v_seq = np.zeros((tot_frm, hh, ww), dtype=np.uint8)

    # read data
    with open(seq_path, 'rb') as fp:
        for i in range(tot_frm):
            fp.seek(int(blk_size * (start_frm + i)), 0)  # skip frames
            y_frm = np.fromfile(fp, dtype=np.uint8, count=y_size).reshape(h, w)
            #y_frm = np.fromfile(fp, dtype=np.uint8, count=y_size)
            #v_frm = np.fromfile(fp, dtype=np.uint8, count=v_size)
            if only_y:
                y_seq[i, ...] = y_frm
            else:
                u_frm = np.fromfile(fp, dtype=np.uint8, \
                                    count=u_size).reshape(hh, ww)
                v_frm = np.fromfile(fp, dtype=np.uint8, \
                                    count=v_size).reshape(hh, ww)
                y_seq[i, ...], u_seq[i, ...], v_seq[i, ...] = y_frm, u_frm, v_frm

    if only_y:
        return y_seq
    else:
        return y_seq, u_seq, v_seq

def write_yuv(raw_write_path, y_seq, u_seq, v_seq, h, w, tot_frm, yuv_type='420p', start_frm=0, only_y=True):
    if yuv_type == '420p':
        hh, ww = h // 2, w // 2
    elif yuv_type == '444p':
        hh, ww = h, w
    else:
        raise Exception('yuv_type not supported.')

    y_size, u_size, v_size = h * w, hh * ww, hh * ww
    blk_size = y_size + u_size + v_size


    # 写入
    y_seq_write = np.reshape(y_seq, (h * w))
    u_seq_write = np.reshape(u_seq, (h * w) // 4)
    v_seq_write = np.reshape(v_seq, (h * w) // 4)

    yuv_seq = np.concatenate((y_seq_write, u_seq_write, v_seq_write), axis=0)
    with open(raw_write_path, 'wb') as fp:
        for i in range(tot_frm):
            fp.seek(int(blk_size * (start_frm + i)), 0)  # skip frames
            yuv_seq.tofile(fp)



# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    raw_y, raw_u, raw_v = import_yuv(
        seq_path=raw_yuv_path, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=False
        )
    write_yuv(raw_write_path=raw_write_path, y_seq=raw_y, u_seq=raw_u, v_seq=raw_v, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=False)
    gt_image = raw_y[[0], :, :]
    gt_image = np.reshape(gt_image, (1080, 1920))
    im = Image.fromarray(gt_image)
    im.save("out.jpeg")

#提取多帧
import numpy as np
from PIL import Image
raw_yuv_path = r'C:\Users\weiwei\Desktop\0.yuv'
raw_write_path = r'C:\Users\weiwei\Desktop\1.yuv'
h, w, nfs = 240, 416, 50
def import_yuv(seq_path, h, w, tot_frm, yuv_type='420p', start_frm=0, only_y=True):
    if yuv_type == '420p':
        hh, ww = h // 2, w // 2
    elif yuv_type == '444p':
        hh, ww = h, w
    else:
        raise Exception('yuv_type not supported.')

    y_size, u_size, v_size = h * w, hh * ww, hh * ww
    blk_size = y_size + u_size + v_size

    # init
    y_seq = np.zeros((tot_frm, h, w), dtype=np.uint8)
    if not only_y:
        u_seq = np.zeros((tot_frm, hh, ww), dtype=np.uint8)
        v_seq = np.zeros((tot_frm, hh, ww), dtype=np.uint8)

    # read data
    with open(seq_path, 'rb') as fp:
        for i in range(tot_frm):
            fp.seek(int(blk_size * (start_frm + i)), 0)  # skip frames
            y_frm = np.fromfile(fp, dtype=np.uint8, count=y_size).reshape(h, w)
            #y_frm = np.fromfile(fp, dtype=np.uint8, count=y_size)
            #v_frm = np.fromfile(fp, dtype=np.uint8, count=v_size)
            if only_y:
                y_seq[i, ...] = y_frm
            else:
                u_frm = np.fromfile(fp, dtype=np.uint8, \
                                    count=u_size).reshape(hh, ww)
                v_frm = np.fromfile(fp, dtype=np.uint8, \
                                    count=v_size).reshape(hh, ww)
                y_seq[i, ...], u_seq[i, ...], v_seq[i, ...] = y_frm, u_frm, v_frm

    if only_y:
        return y_seq
    else:
        return y_seq, u_seq, v_seq

def write_yuv(raw_write_path, y_seq, u_seq, v_seq, h, w, tot_frm, yuv_type='420p', start_frm=0, only_y=True):
    if yuv_type == '420p':
        hh, ww = h // 2, w // 2
    elif yuv_type == '444p':
        hh, ww = h, w
    else:
        raise Exception('yuv_type not supported.')

    y_size, u_size, v_size = h * w, hh * ww, hh * ww
    blk_size = y_size + u_size + v_size


    # 写入
    y_seq_write = np.reshape(y_seq, (tot_frm, h * w))
    u_seq_write = np.reshape(u_seq, (tot_frm, h * w // 4))
    v_seq_write = np.reshape(v_seq, (tot_frm, h * w // 4))

    #yuv_seq = np.concatenate((y_seq_write, u_seq_write, v_seq_write), axis=0)
    with open(raw_write_path, 'wb') as fp:
        for i in range(tot_frm):
            fp.seek(int(blk_size * (start_frm + i)), 0)  # skip frames
            y_seq_write[i].tofile(fp)
            u_seq_write[i].tofile(fp)
            v_seq_write[i].tofile(fp)



# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    raw_y, raw_u, raw_v = import_yuv(
        seq_path=raw_yuv_path, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=False
        )
    write_yuv(raw_write_path=raw_write_path, y_seq=raw_y, u_seq=raw_u, v_seq=raw_v, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=False)
    gt_image = raw_y[[0], :, :]
    gt_image = np.reshape(gt_image, (1080, 1920))
    im = Image.fromarray(gt_image)
    im.save("out.jpeg")