蟒蛇 | Numpy 数组图像转换

Question

我有一个形状为(1000, 50, 100, 3) (class 'numpy.ndarray') 的 Numpy 图像数组，其中包含 1000 个图像 RGB（高度 = 50，宽度 = 100，通道 = 3）。我想首先将RGB 值转换为YUV 值并重新调整它们以获得yuv 值。下面给出了逐像素转换器的原型实现。

我的问题：有没有一种简单的方法可以进行这种转换？

def yuv(_pixel):
    R, G, B = _pixel[0], _pixel[1], _pixel[2]
    Y = 0.299 *  R + 0.587 * G + 0.114 * B
    y = Y / 127.5 - 1
    u = (0.493 * (B - Y)) / 127.5 - 1
    v = (0.887 * (R - Y)) / 127.5 - 1
    return np.array([y, u, v])

Answer 1

你看过numpy.apply_along_axis吗？

你可以这样做：

images_yuv = np.apply_along_axis( yuv, -1, images_rgb)

编辑：混淆了参数的顺序

Answer 2

您可以对转换进行矢量化处理，以便同时转换所有 R、G 和 B 像素：

def yuv_vec(images):
    R, G, B = images[:, :, :, 0], images[:, :, :, 1], images[:, :, :, 2]
    y = (0.299 *  R + 0.587 * G + 0.114 * B) / 127.5 - 1
    u = (0.493 * (B - y)) / 127.5 - 1
    v = (0.887 * (R - y)) / 127.5 - 1
    yuv_img = np.empty(images.shape)
    yuv_img[:, :, :, 0] = y
    yuv_img[:, :, :, 1] = u
    yuv_img[:, :, :, 2] = v
    return yuv_img

为了计时性能，我将展示问题中所示 yuv 函数的简短嵌套循环实现：

def yuv(_pixel):
    R, G, B = _pixel[0], _pixel[1], _pixel[2]
    y = (0.299 *  R + 0.587 * G + 0.114 * B) / 127.5 - 1
    u = (0.493 * (B - Y)) / 127.5 - 1
    v = (0.887 * (R - Y)) / 127.5 - 1
    return np.array([y, u, v])

def yuvloop(imgs):
    yuvimg = np.empty(imgs.shape)
    for n in range(imgs.shape[0]):
        for i in range(imgs.shape[1]):
            for j in range(imgs.shape[2]):
                yuvimg[n, i, j] = yuv(imgs[n, i, j])
    return yuvimg

一些速度比较：

imgs = np.random.randint(0, 256, size=(100, 50, 100, 3))
%timeit yuvloop(imgs)
# Out: 8.79 s ± 265 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
% timeit np.apply_along_axis(yuv, -1, imgs)
# Out: 9.92 s ± 360 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit yuv_vec(imgs)
# Out: 34.4 ms ± 385 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

所以这比在像素上循环快 256 倍。使用np.apply_along_axis 似乎更慢。三者的结果都是一样的。
我将测试样本的大小减少到 100 张图像，否则测试会太慢。