1. 程式人生 > 程式設計 >利用python讀取YUV檔案 轉RGB 8bit/10bit通用

利用python讀取YUV檔案 轉RGB 8bit/10bit通用

注:本文所指的YUV均為YUV420中的I420格式(最常見的一種),其他格式不能用以下的程式碼。

位深為8bit時,每個畫素佔用1位元組,對應檔案指標的fp.read(1);

位深為10bit時,每個畫素佔用2位元組,對應檔案指標的fp.read(2);

然後使用 int.from_bytes() 方法將二進位制轉換為int型數字。

以下程式可以讀8bit或10bit位深的YUV,需要指定從第幾幀開始讀、一共讀多少幀。

它返回三個陣列,其shape分別為:Y [frame,W,H] U [frame,W/2,H/2] V [frame,H/2]

當只讀1幀時它返回:Y [W,H] U [W/2,H/2] V [W/2,H/2]

# -*- coding: utf-8 -*-
 
import math
from functools import partial
import numpy as np
import matplotlib.pyplot as plt
 
 
def readyuv420(filename,bitdepth,H,startframe,totalframe,show=False):
  # 從第startframe(含)開始讀(0-based),共讀totalframe幀
 
  uv_H = H // 2
  uv_W = W // 2
 
  if bitdepth == 8:
    Y = np.zeros((totalframe,W),np.uint8)
    U = np.zeros((totalframe,uv_H,uv_W),np.uint8)
    V = np.zeros((totalframe,np.uint8)
  elif bitdepth == 10:
    Y = np.zeros((totalframe,np.uint16)
    U = np.zeros((totalframe,np.uint16)
    V = np.zeros((totalframe,np.uint16)
 
  plt.ion()
 
  bytes2num = partial(int.from_bytes,byteorder='little',signed=False)
 
  bytesPerPixel = math.ceil(bitdepth / 8)
  seekPixels = startframe * H * W * 3 // 2
  fp = open(filename,'rb')
  fp.seek(bytesPerPixel * seekPixels)
 
  for i in range(totalframe):
 
    for m in range(H):
      for n in range(W):
        if bitdepth == 8:
          pel = bytes2num(fp.read(1))
          Y[i,m,n] = np.uint8(pel)
        elif bitdepth == 10:
          pel = bytes2num(fp.read(2))
          Y[i,n] = np.uint16(pel)
 
    for m in range(uv_H):
      for n in range(uv_W):
        if bitdepth == 8:
          pel = bytes2num(fp.read(1))
          U[i,n] = np.uint8(pel)
        elif bitdepth == 10:
          pel = bytes2num(fp.read(2))
          U[i,n] = np.uint16(pel)
 
    for m in range(uv_H):
      for n in range(uv_W):
        if bitdepth == 8:
          pel = bytes2num(fp.read(1))
          V[i,n] = np.uint8(pel)
        elif bitdepth == 10:
          pel = bytes2num(fp.read(2))
          V[i,n] = np.uint16(pel)
 
    if show:
      print(i)
      plt.subplot(131)
      plt.imshow(Y[i,:,:],cmap='gray')
      plt.subplot(132)
      plt.imshow(U[i,cmap='gray')
      plt.subplot(133)
      plt.imshow(V[i,cmap='gray')
      plt.show()
      plt.pause(1)
      #plt.pause(0.001)
 
  if totalframe==1:
    return Y[0],U[0],V[0]
  else:
    return Y,U,V
 
 
if __name__ == '__main__':
  #y,u,v = readyuv420(r'F:\_commondata\video\176x144 qcif\football_qcif.yuv',8,176,144,1,5,True)
  y,v = readyuv420(r'F:\_commondata\video\1920x1080 B\RitualDance_1920x1080_60fps_10bit_420.yuv',10,1920,1080,True)
  print(y.shape,u.shape,v.shape)

以下程式將YUV轉為RGB(只能讀8bit位深的YUV),返回1個數組,其shape為: [frame,3]

# -*- coding: utf-8 -*-
import cv2
import numpy as np
import matplotlib.pyplot as plt
 
 
def yuv2rgb(yuvfilename,show=False,out=False):
  # 從第startframe(含)開始讀(0-based),共讀totalframe幀
  arr = np.zeros((totalframe,3),np.uint8)
  
  plt.ion()
  with open(yuvfilename,'rb') as fp:
    seekPixels = startframe * H * W * 3 // 2
    fp.seek(8 * seekPixels) #跳過前startframe幀
    for i in range(totalframe):
      print(i)
      oneframe_I420 = np.zeros((H*3//2,np.uint8)
      for j in range(H*3//2):
        for k in range(W):
          oneframe_I420[j,k] = int.from_bytes(fp.read(1),signed=False)
      oneframe_RGB = cv2.cvtColor(oneframe_I420,cv2.COLOR_YUV2RGB_I420)
      if show:
        plt.imshow(oneframe_RGB)
        plt.show()
        plt.pause(0.001)
      if out:
        outname = yuvfilename[:-4]+'_'+str(startframe+i)+'.png'
        cv2.imwrite(outname,oneframe_RGB[:,::-1])
      arr[i] = oneframe_RGB
  return arr
 
if __name__ == '__main__':
  video = yuv2rgb(r'D:\_workspace\akiyo_qcif.yuv',False,True)

用ffmpeg也可以,比如你需要將yuv的第8幀輸出成一個png:

ffmpeg -s 176x144 -i akiyo_qcif.yuv -filter:v select="between(n\,8\,8)" out.png

以上這篇利用python讀取YUV檔案 轉RGB 8bit/10bit通用就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。