Search code examples
pythonnumpyyuv

Import YUV as a byte array


I'm working on a project where I have to apply threshold to a YUV420_SP_NV21 image (taken from an Android camera) to determine which pixels are 'black' and which are 'white'.

Therefore, I want to import it in Python as a bytearray (using OpenCV, NumPy, PIL, ...), so I can do some quick bitwise operations on the Y-values.



However, when I try to import the image using the following methods, I get useless outputs:

When I type

import cv2
import numpy as np

img = cv2.imread('test.yuv')
imgArr = np.array(img)

print(img)
print(imgArr)

I get as output:

None
None



And when I type

import numpy as np

f = open('test.yuv', 'rb')
img = f.read()
imgArr = np.array(img)

I get some useless character sequence.
And when I type now (for exemple)

print(imgArr[0])

I get as output:

IndexError: too many indices for array

Which means imgArr is no array at all!



Can anyone please tell me what I'm doing wrong? Thanks in advance!


Solution

  • You really should provide a sample image so folks can help you better - else they spend ages developing code for what they guess is the image format you mean and it is a waste of effort if they guess wrong.

    Assuming your data look like this from Wikipedia, you can see that all the Y samples come first, followed by all the U samples and all the V samples. However there are only 1/4 as many U/V samples as Y samples because these two components are subsampled at 2:1 ratio both horizontally and vertically:

    enter image description here

    So, the idea is to read in the entire set of YUV samples, then take the first w*h bytes as the Y values, the next w*h/4 samples as U and the next w*h/4 samples as V:

    import numpy as np
    
    # Guess a width and height and derive number of pixels in image
    w,h = 2048,1536
    px = w*h
    
    # Read entire file into YUV
    YUV = np.fromfile('NV21_2048x1536.yuv',dtype='uint8')
    
    # Take first h x w samples and reshape as Y channel
    Y = YUV[0:w*h].reshape(h,w)
    
    # Take next px/4 samples as U
    U = YUV[px:(px*5)//4].reshape(h//2,w//2)
    
    # Take next px/4 samples as V
    V = YUV[(px*5)//4:(px*6)//4].reshape(h//2,w//2)
    
    # Undo subsampling of U and V by doubling height and width
    Ufull = U.copy().resize((w,h))
    Vfull = V.copy().resize((w,h))
    

    I have no idea what you plan to do next, so I'll leave it at that for now!


    Keywords: NV21, YUV, YUV420, Android, YUV430P, Pillow, PIL, Python, image, image processing.