如何加速Python中NV12转YUV444操作 替代原生for循环实现高效转换

我最初的NV12转YUV444采用Python原生双层for循环实现，代码如下：

for h in range(self.img_shape[0]):
    # 为推理框架规整YUV444数据
    for w in range(self.img_shape[1]):
        yuv444_res[h][w][0] = (nv12_y_data[h * self.img_shape[1] +w]).astype(np.int8)
        yuv444_res[h][w][1] = (nv12_u_data[int(h / 2) * int(self.img_shape[1] / 2) +int(w / 2)]).astype(np.int8)
        yuv444_res[h][w][2] = (nv12_v_data[int(h / 2) * int(self.img_shape[1] / 2) +int(w / 2)]).astype(np.int8)

由于Python原生for循环运行速度远低于NumPy向量化运算，我希望找到基于NumPy的优化实现方案。

2021年6月15日更新 #

我基于花式索引实现了优化版本，参考公开技术资料编写，代码如下：

yuv444 = np.empty([self.height, self.width, 3], dtype=np.uint8)
yuv444[:, :, 0] = nv12_data[:self.width * self.height].reshape(
    self.height, self.width)
u = nv12_data[self.width * self.height::2].reshape(
    self.height // 2, self.width // 2)
yuv444[:, :, 1] = Image.fromarray(u).resize((self.width, self.height))
v = nv12_data[self.width * self.height + 1::2].reshape(
    self.height // 2, self.width // 2)
yuv444[:, :, 2] = Image.fromarray(v).resize((self.width, self.height))

data[0] = yuv444.astype(np.int8)

使用PIL的resize方法替代已废弃的imresize接口后，该代码运行结果与原有for循环实现100%匹配。

2021年6月19日更新 #

参考Rotem提供的实现方案后，我找到了运行速度更快的OpenCV实现版本，代码如下：

# nv12_data已被重塑为一维数组
y = nv12_data[:self.width * self.height].reshape(
    self.height, self.width)
shrunk_u = nv12_data[self.width * self.height::2].reshape(
    self.height // 2, self.width // 2)
shrunk_v = nv12_data[self.width * self.height + 1::2].reshape(
    self.height // 2, self.width // 2)
u = cv2.resize(shrunk_u, (self.width, self.height),
               interpolation=cv2.INTER_NEAREST)
v = cv2.resize(shrunk_v, (self.width, self.height),
               interpolation=cv2.INTER_NEAREST)
yuv444 = np.dstack((y, u, v))

我对1000张图片的处理进行了耗时对比，结果显示OpenCV方案速度更快，且运算结果和原生实现完全一致：

cv time: 4.417593002319336, pil time: 5.395732164382935

2021年6月25日更新 #

经测试发现，Pillow的resize方法在不同版本中的默认resample参数值存在差异：

• 5.1.0版本：

def resize(self, size, resample=NEAREST, box=None):

• 8.1.0版本：

def resize(self, size, resample=BICUBIC, box=None, reducing_gap=None):

因此调用resize接口时最好显式指定重采样策略，避免不同版本运行结果不一致的问题。

内容的提问来源于stack exchange，提问作者Johnzy