ffmpeg配合python实现RGB到YUV444的转换

要通过python将RGB转成YUV444。

opencv提供了RGB转YUV的工具，如果不在乎是哪一种YUV采样方式，直接使用以下的方式即可：

import cv2

img = cv2.imread(img_path)
img_YUV= cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
img_YCrCb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)

opencv中的YUV是YUV420，并不支持YUV444。YUV420和YCrCb的差异，可以参考这篇文章

了解到ffmpeg支持python，且可以实现RGB到YUV444的转换，本文主要介绍一下整个过程如何实现。

安装ffmpeg

sudo apt-get install ffmpeg

安装ffmpy

pip install ffmpy

在python中使用ffmpeg，其实就是讲ffmpeg指令放入ffmpy中， 并在命令行中执行代码。

说明和使用样例可以参见以下链接。
ffmpy文档

在pythong中使用ffmpeg实现RGB到YUV444的转换

首先在命令行中输入

ffmpeg -pix_fmts

可以查看ffmpeg所有支持的格式，这边仅读取部分示意。

FLAGS NAME            NB_COMPONENTS BITS_PER_PIXEL
-----
IO... yuv420p                3            12
IO... yuyv422                3            16
IO... rgb24                  3            24
IO... bgr24                  3            24
IO... yuv422p                3            16
IO... yuv444p                3            24
IO... yuv410p                3             9
IO... yuv411p                3            12
IO... gray                   1             8
IO..B monow                  1             1
IO..B monob                  1             1
I..P. pal8                   1             8
IO... yuvj420p               3            12
IO... yuvj422p               3            16
IO... yuvj444p               3            24
IO... uyvy422                3            16
..... uyyvyy411              3            12
IO... bgr8                   3             8
.O..B bgr4                   3             4
IO... bgr4_byte              3             4
IO... rgb8                   3             8
.O..B rgb4                   3             4
IO... rgb4_byte              3             4
IO... nv12                   3            12
IO... nv21                   3            12
IO... argb                   4            32
IO... rgba                   4            32
IO... abgr                   4            32
IO... bgra                   4            32
...

RGB到YUV的转换

下面定义的函数可以实现RGB到YUV的转换

import os
import cv2
from ffmpy import FFmpeg

def convert(img_path,H,W,output_path):
	size = '{}x{}'.format(W,H)
	ff = FFmpeg(inputs={img_path,None},
				outputs={output_path:'-s {} -pix_fmt yuv444p'.format(size)})
	print(ff.cmd)
	ff.run()

def main():
	img_bgr = cv2.imread(img_path)
	H,W = img_bgr[:2]
	convert(img_path,H,W,'test.yuv')#执行后在指定路径中将生成test.yuv文件

if __name__ == '__main__':
	main()

再执行过一次后，由于指定路径已存在test.yuv文件，再次执行时，汇报exited with status 1的错。更改路径、文件名或者将路径中的同名文件删除即可解决。

python读取YUV文件

由于ffmpeg执行后的结果是保存一个yuv文件，我们还需要一个函数来读取yuv文件供后续使用。

import numpy as np
from matplotlib import pyplot as plt

def yuv_read(filename,dims,numfrm,startfrm):  
    fp=open(filename,'rb')  
    blk_size = int(np.prod(dims) *3/2)
    fp.seek(blk_size*startfrm,0)  
    Y=[]  
    U=[]  
    V=[]  
    print(dims[0])  
    print(dims[1]) 
    d00=dims[0]  #如果是YUV420，此处为 d00=dims[0] //2
    d01=dims[1]   #如果是YUV420，此处为 d01=dims[1] //2
    print(d00)
    print(d01) 
    Yt=np.zeros((dims[0],dims[1]),np.uint8,'C')  
    Ut=np.zeros((d00,d01),np.uint8,'C')  
    Vt=np.zeros((d00,d01),np.uint8,'C')  
    for i in range(numfrm):  
        for m in range(dims[0]):  
            for n in range(dims[1]):  
                #print m,n  
                Yt[m,n]=ord(fp.read(1))  
        for m in range(d00):  
            for n in range(d01):  
                Ut[m,n]=ord(fp.read(1))  
        for m in range(d00):  
            for n in range(d01):  
                Vt[m,n]=ord(fp.read(1))  
        Y=Y+[Yt]  
        U=U+[Ut]  
        V=V+[Vt]  
    fp.close()  
    return (Y,U,V)

def main():
	data = yuv_read('test.yuv',(H,W),1,0)
	# 查看Y\U\V分量
	plt.figure()
	plt.subplt(1,3,1)
	plt.imshow(data[0][0])# Y分量
	plt.subplt(1,3,2)
	plt.imshow(data[1][0])# U分量
	plt.subplt(1,3,3)
	plt.imshow(data[2][0])# V分量
	plt.show