采用Python实现快速YUV转RGB

目标

采用Python脚本实现快速的YUV图像二进制(BIN)文件到sRGB-24bit图像的转换，并保存为PNG文件。

解决方法

一般来说，YUV转RGB的主要手段有三种：

• libYUV, from Google Chromium Project, Open Source.
• FFmpeg, the most popular tool to decode images and videos across all web/pc platforms.(BiliBili也是用的这个库做的视频解码）
• OpenCV, open source project mainly providing all popular and stable algorithms in computer vision.

各自特点：
采用libYUV会更快（一般4倍），采用FFmpeg效果可能更贴近人眼视觉（？），但采用OpenCV最方便，因为libYUV需要从源码根据平台修改编译，并没有库的形式直接调用，FFmpeg虽然有库，但是对python不友好（可能是我个人对基于FFmpeg的库不了解？），opencv就很简单，二进制的Yuv数据读进来就能直接转sRGB，而且和Numpy无缝结合，极为便利。

Show me the code right now!

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


# 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(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_YUV2RGB_NV21)
    elif t == YUV_NV12:
        im_rgb = cv2.cvtColor(im_yuv, cv2.COLOR_YUV2RGB_NV12)
    else:
        print('not implemented yet!')
        assert False
    return im_rgb[:, :w]


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.nv21')
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV(f_, main_w, main_h, pitch, 3, YUV_NV12)
        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('.nv21', '.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.nv21')
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV(f_, aux_w, aux_h, pitch, 3, YUV_NV12)
        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('.nv21', '.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.nv21')
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV(f_, bokeh_w, bokeh_h, pitch, 3, YUV_NV12)
        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('.nv21', '.png')
        cv2.imwrite(fn_ , im_)
    print('================================================================')

整个代码中其实核心的只有两行(或2个函数调用）：

1. np.frombuffer()：从二进制字节流中结构化为numpy数组，如果字节流来自网络则选择大端的>{size}B,如果不是uint8，而是float类型，则选择<{size}f，这个规则是numpy适应struct库得到的。
2. cv2.cvtColor()：opencv从numpy数组读取图像，然后转化到指定的颜色空间如sRGB.