采用Python实现快速YUV转RGB

目标

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

解决方法

一般来说,YUV转RGB的主要手段有三种:

  • libYUV, from Google Chromium Project, Open Source.
  • FFmpeg, the most popular tools to decode images and video 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


def DecodeYUV(fp, w, h, c, t): # file pointer, width, height, channel, YUV type
    if c == 3:
        size_ = w*h*3//2
    else:
        size_ = w*h
    data = f_.read(size_)
    if c == 3:
        im_yuv = np.frombuffer(data, dtype='<'+str(size_)+'B').reshape([h*3//2, main_w])
    else:
        im_yuv = np.frombuffer(data, dtype='<'+str(size_)+'B').reshape([h, w])
    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


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 = 3264
    main_h = 2448
    files = glob.glob(yuv_dir + '/*_input*_3264x2448.yuv')
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV(f_, main_w, main_h, 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('.yuv', '.png')
        cv2.imwrite(fn_ , im_)
    
    print('****** PROCESSING AUX ******')
    # find aux images
    aux_w = 1600
    aux_h = 1200
    files = glob.glob(yuv_dir + '/*_1600x1200.yuv')
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV(f_, aux_w, aux_h, 1, 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('.yuv', '.png')
        cv2.imwrite(fn_ , im_)
    
    print('****** PROCESSING BOKEH ******')
    # find aux images
    bokeh_w = main_w
    bokeh_h = main_h
    files = glob.glob(yuv_dir + '/*_out_*.yuv')
    focus_file = open('focus.txt', 'wt')
    for i in range(len(files)):
        print(files[i])
        f_ = open(files[i], 'rb')
        im_ = DecodeYUV(f_, bokeh_w, bokeh_h, 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('.yuv', '.png')
        cv2.imwrite(fn_ , im_)
        # save the focus information from filename
        beg_id = files[i].rfind('Focus')
        end_id = files[i].rfind('_levle')
        focus_str = files[i][beg_id + 5:end_id]
        if i==0:
            focus_file.write(focus_str)
        else:
            focus_file.write('\n' + focus_str)
    focus_file.close()
    print('================================================================')
上一篇:rgb转16进制js方法,npm插件


下一篇:【546】灰度图array转为RGB三通道array