采用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.