采用Numpy实现准确的YUV转sRGB

目标

利用Python利用标准公式实现准确的YUV转RGB（从YCrCb空间转到sRGB）

解决

1.Numpy的处理速度比原生python的for循环更快，且支持跨平台执行；
2.利用Numpy实现将YUV格式二进制流转到sRGB格式，并保存为PNG图像；

代码

import numpy as np
import cv2
import sys
import os
import platform
import glob

def Windows():
    return platform.system() == "Windows"


YUV_NV21 = 0
YUV_NV12 = 1

postfix = '.yuv'

# param:
# @f: file pointer
# @w: image width
# @h: image height
# @p: pitch size
# @c: image color channel
# @t YUV type, NV21 or NV12
def DecodeYUV_OpenCV(f, w, h, p, c, t):
    if c == 3:
        size_ = p*h*3//2
    else:
        size_ = p*h
    data = f.read(size_)
    if c == 3:
        im_yuv = np.frombuffer(data, dtype='<'+str(size_)+'B').reshape([h*3//2, p])
    else:
        im_yuv = np.frombuffer(data, dtype='<'+str(size_)+'B').reshape([h, p])
    if c==1:
        return im_yuv
    im_rgb = None
    if t == YUV_NV21:
        im_rgb = cv2.cvtColor(im_yuv, cv2.COLOR_YUV2BGR_NV21)
    elif t == YUV_NV12:
        im_rgb = cv2.cvtColor(im_yuv, cv2.COLOR_YUV2BGR_NV12)
    else:
        print('not implemented yet!')
        assert False
    return im_rgb[:, :w]
    
    
def DecodeYUV_Numpy(f, w, h, p, c, t):
    assert w % 2 == 0
    assert h % 2 == 0
    if c == 3:
        size_ = p*h*3//2
    else:
        size_ = p*h
    data = f.read(size_)
    if c == 3:
        im_yuv = np.frombuffer(data, dtype='<'+str(size_)+'B').reshape([h*3//2, p])
    else:
        im_yuv = np.frombuffer(data, dtype='<'+str(size_)+'B').reshape([h, p])
    if c==1:
        return im_yuv
    # directly use the equation to convert from YUV to RGB
    im_yuv = np.float32(im_yuv)
    y_ = im_yuv[:h, :w]
    uv_ = im_yuv[h:, :w]
    im_rgb = None
    u_full = np.zeros([h, w], np.float32)
    v_full = np.zeros([h, w], np.float32)
    if t == YUV_NV12: # yyyyuvuv
        u_ = uv_[:, 0:w:2]
        v_ = uv_[:, 1:w:2]
    elif t == YUV_NV21: # yyyyvuvu
        u_ = uv_[:, 1:w:2]
        v_ = uv_[:, 0:w:2]    
    else:
        print("Unsupported YUV format!")
        assert False
    #u_full = cv2.resize(u_, (w, h), cv2.INTER_LINEAR)
    #v_full = cv2.resize(v_, (w, h), cv2.INTER_LINEAR)
    u_full[0:h:2, 0:w:2] = u_
    u_full[0:h:2, 1:w:2] = u_
    u_full[1:h:2, 0:w:2] = u_
    u_full[1:h:2, 1:w:2] = u_
    
    v_full[0:h:2, 0:w:2] = v_
    v_full[0:h:2, 1:w:2] = v_
    v_full[1:h:2, 0:w:2] = v_
    v_full[1:h:2, 1:w:2] = v_
    
    u_full -= 128
    v_full -= 128
    r_ = np.maximum(0, np.minimum(255, y_ + 1.402*v_full))
    g_ = np.maximum(0, np.minimum(255, y_ - 0.344*u_full - 0.714*v_full))
    b_ = np.maximum(0, np.minimum(255, y_ + 1.772*u_full))
    
    return np.uint8(np.stack([b_, g_, r_], axis=2))


if __name__ == '__main__':
    assert len(sys.argv) == 2
    print('================ CONVERT YUV 2 RGB FOR PRAGUE ONLY =============')
    # create dirs for output images
    dir_main = 'picMain'
    dir_aux = 'picAux'
    dir_bokeh = 'Bokeh'
    cmd_mkdir = ''
    if Windows():
        cmd_mkdir = 'md '
    else:
        cmd_mkdir = 'mkdir -p '
    os.system(cmd_mkdir + dir_main)
    os.system(cmd_mkdir + dir_aux)
    os.system(cmd_mkdir + dir_bokeh)
    
    yuv_dir = sys.argv[1]
    print('YUV DIR : ' + yuv_dir)
    print('****** PROCESSING MAIN ******')
    # find main images
    main_w = 4608
    main_h = 3456
    pitch = main_w
    files = glob.glob(yuv_dir + '/*_BokehInput0'+postfix)
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV_Numpy(f_, main_w, main_h, pitch, 3, YUV_NV21)
        f_.close()
        # save RGB images
        fn_start = ''
        if Windows():
            fn_start = files[i].rfind('\\')
        else:
            fn_start = files[i].rfind('/')
        fn_ = dir_main + files[i][fn_start:]
        fn_ = fn_.replace(postfix, '.png')
        cv2.imwrite(fn_ , im_)
    
    print('****** PROCESSING AUX ******')
    # find aux images
    aux_w = 2592
    aux_h = 1944
    pitch = 3072
    files = glob.glob(yuv_dir + '/*_BokehInput1'+postfix)
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV_Numpy(f_, aux_w, aux_h, pitch, 3, YUV_NV21)
        f_.close()
        # save RGB images
        fn_start = ''
        if Windows():
            fn_start = files[i].rfind('\\')
        else:
            fn_start = files[i].rfind('/')
        fn_ = dir_aux + files[i][fn_start:]
        fn_ = fn_.replace(postfix, '.png')
        cv2.imwrite(fn_ , im_)
    
    print('****** PROCESSING BOKEH ******')
    # find aux images
    bokeh_w = main_w
    bokeh_h = main_h
    pitch = bokeh_w
    files = glob.glob(yuv_dir + '/*_BokehOutput'+postfix)
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV_Numpy(f_, bokeh_w, bokeh_h, pitch, 3, YUV_NV21)
        f_.close()
        # save RGB images
        fn_start = ''
        if Windows():
            fn_start = files[i].rfind('\\')
        else:
            fn_start = files[i].rfind('/')
        fn_ = dir_bokeh + files[i][fn_start:]
        fn_ = fn_.replace(postfix, '.png')
        cv2.imwrite(fn_ , im_)
    print('================================================================')

其中DecodeYUV_Numpy()函数就是最终解决方案：注释掉的一种更为准确的转换方法是将复制上采样改为双线/立方插值上采样UV通道，但这种方法要慢一些。