yuv转rgb

yuv转rgb

注意字深和yuv具体的存储类型是NV12 NV21 I420 YV12中的哪种
yuv相关介绍

import cv2
from numpy import *  
 
def yuv_import(filename,dims,numfrm,startfrm,type):  
    fp=open(filename,'rb')  
    if type == '8bit':
        blk_size = prod(dims)*3//2 # Y U V
    elif type == '10bit':
        blk_size = prod(dims)*3//2*2 # Y U V
    fp.seek(blk_size*startfrm,0)  
    Y=[]  
    U=[]  
    V=[]  
    d00=dims[0]//2  
    d01=dims[1]//2  
    Yt=zeros((dims[0],dims[1]),uint8,'C')  
    Ut=zeros((d00,d01),uint8,'C')  
    Vt=zeros((d00,d01),uint8,'C')  
    for i in range(numfrm):  
        for m in range(dims[0]):  
            for n in range(dims[1]):  
                if type == '8bit':
                    Yt[m,n] = ord(fp.read(1))
                elif type == '10bit':
                    Yt[m,n] = (ord(fp.read(1)) + ord(fp.read(1))*255)//4 
        for m in range(d00):  
            for n in range(d01):  
                if type == '8bit':
                    Ut[m,n] = ord(fp.read(1))
                elif type == '10bit':
                    Ut[m,n] = (ord(fp.read(1)) + ord(fp.read(1))*255)//4 
        for m in range(d00):  
            for n in range(d01):  
                if type == '8bit':
                    Vt[m,n] = ord(fp.read(1))
                elif type == '10bit':
                    Vt[m,n] = (ord(fp.read(1)) + ord(fp.read(1))*255)//4 
        Y=Y+[Yt]  
        U=U+[Ut]  
        V=V+[Vt]  
    fp.close()  
    return (Y,U,V)  

def yuv2bgr(filename,height,width,numfrm,startfrm,type):  
    fp=open(filename,'rb')  
    if type == '8bit':
        blk_size = height*width*3//2
    elif type == '10bit':
        blk_size = height*width*3//2*2
    fp.seek(blk_size*startfrm,0)  
    video = []
    Yt=zeros((height,width),uint8,'C')  
    Ut=zeros((height//2,width//2),uint8,'C')  
    Vt=zeros((height//2,width//2),uint8,'C')  
    for i in range(numfrm):
        for m in range(height):  
            for n in range(width):  
                if type == '8bit':
                    Yt[m,n] = ord(fp.read(1))
                elif type == '10bit':
                    Yt[m,n] = (ord(fp.read(1)) + ord(fp.read(1))*255)//4 
        for m in range(height//2):  
            for n in range(width//2):  
                if type == '8bit':
                    Ut[m,n] = ord(fp.read(1))
                elif type == '10bit':
                    Ut[m,n] = (ord(fp.read(1)) + ord(fp.read(1))*255)//4 
        for m in range(height//2):  
            for n in range(width//2):   
                if type == '8bit':
                    Vt[m,n] = ord(fp.read(1))
                elif type == '10bit':
                    Vt[m,n] = (ord(fp.read(1)) + ord(fp.read(1))*255)//4 
        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）
        bgr_img = cv2.cvtColor(img, cv2.COLOR_YUV2BGR_I420) # 注意 YUV 的存储格式
        video.append(bgr_img)
        print("Extract frame %d " % (i + 1))
    fp.close()  
    return video

if __name__ == '__main__':
    width=3840
    height=2160
    type = '10bit'

    video = yuv2bgr('dataset/video/SunBath_3840x2160_50fps_10bit.yuv',height,width,10,0,type)  
    video = np.array(video)
    print(video.shape)
    cv2.imwrite('dataset/output/video.jpg',video[0])