当前位置 主页 > 网站技术 > 代码类 >

    利用python读取YUV文件 转RGB 8bit/10bit通用

    栏目:代码类 时间:2019-12-09 18:07

    注:本文所指的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,W/2,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, W, H, startframe, totalframe, show=False):
      # 从第startframe(含)开始读(0-based),共读totalframe帧
     
      uv_H = H // 2
      uv_W = W // 2
     
      if bitdepth == 8:
        Y = np.zeros((totalframe, H, W), np.uint8)
        U = np.zeros((totalframe, uv_H, uv_W), np.uint8)
        V = np.zeros((totalframe, uv_H, uv_W), np.uint8)
      elif bitdepth == 10:
        Y = np.zeros((totalframe, H, W), np.uint16)
        U = np.zeros((totalframe, uv_H, uv_W), np.uint16)
        V = np.zeros((totalframe, uv_H, uv_W), 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, m, 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, m, n] = np.uint8(pel)
            elif bitdepth == 10:
              pel = bytes2num(fp.read(2))
              U[i, m, 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, m, n] = np.uint8(pel)
            elif bitdepth == 10:
              pel = bytes2num(fp.read(2))
              V[i, m, 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, u, v = readyuv420(r'F:\_commondata\video\1920x1080 B\RitualDance_1920x1080_60fps_10bit_420.yuv', 10, 1920, 1080, 0, 5, True)
      print(y.shape,u.shape,v.shape)
    

    以下程序将YUV转为RGB(只能读8bit位深的YUV),返回1个数组,其shape为: [frame,W,H,3]

    # -*- coding: utf-8 -*-
    import cv2
    import numpy as np
    import matplotlib.pyplot as plt
     
     
    def yuv2rgb(yuvfilename, W, H, startframe, totalframe, show=False, out=False):
      # 从第startframe(含)开始读(0-based),共读totalframe帧
      arr = np.zeros((totalframe,H,W,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,W),np.uint8)
          for j in range(H*3//2):
            for k in range(W):
              oneframe_I420[j,k] = int.from_bytes(fp.read(1), byteorder='little', 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', 176, 144, 0, 10, False, True)