NV12 to YUV444 speed up

You may use the process described in my following post, in reverse order (without the RGB part).

Illustration:

Start by creating a synthetic sample image in NV12 format, using FFmpeg (command line tool).
The sample image is used for testing.

Executing from Python using subprocess module:

import subprocess as sp
import shlex

sp.run(shlex.split('ffmpeg -y -f lavfi -i testsrc=size=192x108:rate=1:duration=1 -vcodec rawvideo -pix_fmt nv12 nv12.yuv'))
sp.run(shlex.split('ffmpeg -y -f rawvideo -video_size 192x162 -pixel_format gray -i nv12.yuv -pix_fmt gray nv12_gray.png'))

Read the sample image, and executing the code from your post (used as reference):

import numpy as np
import cv2

nv12 = cv2.imread('nv12_gray.png', cv2.IMREAD_GRAYSCALE)
cols, rows = nv12.shape[1], nv12.shape[0]*2//3

# Reference implementation - using for-loops (the solution is in the part below):
################################################################################
nv12_y_data = nv12[0:rows, :].flatten()
nv12_u_data = nv12[rows:, 0::2].flatten()
nv12_v_data = nv12[rows:, 1::2].flatten()

yuv444_res = np.zeros((rows, cols, 3), np.uint8)

for h in range(rows):
    # centralize yuv 444 data for inference framework
    for w in range(cols):
        yuv444_res[h][w][0] = (nv12_y_data[h * cols + w]).astype(np.int8)
        yuv444_res[h][w][1] = (nv12_u_data[int(h / 2) * int(cols / 2) + int(w / 2)]).astype(np.int8)
        yuv444_res[h][w][2] = (nv12_v_data[int(h / 2) * int(cols / 2) + int(w / 2)]).astype(np.int8)

################################################################################

My suggested solution applies the following stages:

• Separate U and V into two "half size" matrices shrunk_u and shrunk_v.
• Resize shrunk_u and shrunk_v to full image size matrices using cv2.resize.
  In my code sample I used nearest neighbor interpolation for getting the same result as your result.
  It is recommended to replace it with linear interpolation for better quality.
• Use np.dstack for merging Y, U and V into YUV (3 color channels) image.

Here is the complete code sample:

import numpy as np
import subprocess as sp
import shlex
import cv2

sp.run(shlex.split('ffmpeg -y -f lavfi -i testsrc=size=192x108:rate=1:duration=1 -vcodec rawvideo -pix_fmt nv12 nv12.yuv'))
sp.run(shlex.split('ffmpeg -y -f rawvideo -video_size 192x162 -pixel_format gray -i nv12.yuv -pix_fmt gray nv12_gray.png'))
#sp.run(shlex.split('ffmpeg -y -f rawvideo -video_size 192x108 -pixel_format nv12 -i nv12.yuv -vcodec rawvideo -pix_fmt yuv444p yuv444.yuv'))
#sp.run(shlex.split('ffmpeg -y -f rawvideo -video_size 192x324 -pixel_format gray -i yuv444.yuv -pix_fmt gray yuv444_gray.png'))
#sp.run(shlex.split('ffmpeg -y -f rawvideo -video_size 192x108 -pixel_format yuv444p -i yuv444.yuv -pix_fmt rgb24 rgb.png'))
#sp.run(shlex.split('ffmpeg -y -f rawvideo -video_size 192x108 -pixel_format gbrp -i yuv444.yuv -filter_complex "extractplanes=g+b+r[g][b][r],[r][g][b]mergeplanes=0x001020:gbrp[v]" -map "[v]" -vcodec rawvideo -pix_fmt rgb24 yuvyuv.yuv'))
#sp.run(shlex.split('ffmpeg -y -f rawvideo -video#_size 576x108 -pixel_format gray -i yuvyuv.yuv -pix_fmt gray yuvyuv_gray.png'))

nv12 = cv2.imread('nv12_gray.png', cv2.IMREAD_GRAYSCALE)
cols, rows = nv12.shape[1], nv12.shape[0]*2//3

nv12_y_data = nv12[0:rows, :].flatten()
nv12_u_data = nv12[rows:, 0::2].flatten()
nv12_v_data = nv12[rows:, 1::2].flatten()

yuv444_res = np.zeros((rows, cols, 3), np.uint8)

for h in range(rows):
    # centralize yuv 444 data for inference framework
    for w in range(cols):
        yuv444_res[h][w][0] = (nv12_y_data[h * cols + w]).astype(np.int8)
        yuv444_res[h][w][1] = (nv12_u_data[int(h / 2) * int(cols / 2) + int(w / 2)]).astype(np.int8)
        yuv444_res[h][w][2] = (nv12_v_data[int(h / 2) * int(cols / 2) + int(w / 2)]).astype(np.int8)

y = nv12[0:rows, :]
shrunk_u = nv12[rows:, 0::2].copy()
shrunk_v = nv12[rows:, 1::2].copy()

u = cv2.resize(shrunk_u, (cols, rows), interpolation=cv2.INTER_NEAREST)  # Resize U channel (use NEAREST interpolation - fastest, but lowest quality).
v = cv2.resize(shrunk_v, (cols, rows), interpolation=cv2.INTER_NEAREST)  # Resize V channel

yuv444 = np.dstack((y, u, v))

is_eqaul = np.all(yuv444 == yuv444_res)
print('is_eqaul = ' + str(is_eqaul))  # is_eqaul = True

# Convert to RGB for display
yvu = np.dstack((y, v, u))  # Use COLOR_YCrCb2BGR, because it's uses the corrected conversion coefficients.
rgb = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2BGR)

# Show results:
cv2.imshow('nv12', nv12)
cv2.imshow('yuv444_res', yuv444_res)
cv2.imshow('yuv444', yuv444)
cv2.imshow('rgb', rgb)
cv2.waitKey()
cv2.destroyAllWindows()

Input (NV12 displayed as Grayscale):

Output (after converting to RGB):

Seems to be a prime case for fancy indexing (advanced indexing).

Something like this should do the trick, though I didn't verify it on an actual image. I've added a section to reconstruct the image in the beginning, because it is easier to work with the array as a whole than broken into parts. Likely, you can refactor this and avoid splitting it to begin with.

# reconstruct image array
y = nv12_y_data.reshape(self.image_shape[0], self.image_shape[1])
u = nv12_u_data.reshape(self.image_shape[0], self.image_shape[1])
v = nv12_v_data.reshape(self.image_shape[0], self.image_shape[1])
img = np.stack((y,u,v), axis=-1)

# take every index twice until half the range
idx_h = np.repeat(np.arange(img.shape[0] // 2), 2)[:, None]
idx_w = np.repeat(np.arange(img.shape[1] // 2), 2)[None, :]

# convert
yuv444 = np.empty_like(img, dtype=np.uint8)
yuv444[..., 0] = img[..., 0]
yuv444[..., 1] = img[idx_h, idx_w, 1]
yuv444[..., 2] = img[idx_h, idx_w, 2]

If this is along your critical path, and you want to tease out a little more performance, you could consider processing the image channel first, which will be faster on modern CPUs (but not GPUs).

This answer is just another way to do it, and is not the quickest way to get the job done, but definitely should be easy to understand. I have checked the generated files with yuvplayer application as well to confirm it works.

#height mentioned is height of nv12 file and so is the case with width
def convert_nv12toyuv444(filename= 'input.nv12',height=2358,width=2040):

    nv12_data = np.fromfile(filename, dtype=np.uint8)        
    imageSize = (height, width)
    npimg = nv12_data.reshape(imageSize) 
    y_height = npimg.shape[0] * (2/3)
    y_wid = npimg.shape[1]        
    y_height = int(y_height)
    y_wid = int(y_wid)
    y_data= npimg[:y_height,:y_wid]
    uv_data=npimg[y_height:,:y_wid]     
    shrunkU= uv_data[:, 0 : :2] 
    shrunkV= uv_data[:, 1 : :2]       
    u = cv2.resize(shrunkU, (y_wid, y_height),
                   interpolation=cv2.INTER_NEAREST)    
    v = cv2.resize(shrunkV, (y_wid, y_height),
                   interpolation=cv2.INTER_NEAREST)
    yuv444 = np.dstack((y_data, u, v))