python - 使用 NumPy 和 OpenCV 高效地旋转图像并粘贴到更大的图像中

Question

我使用以下代码旋转图像并将它们粘贴到屏幕大小 (4K) 的 Canvas 上，但此实现每旋转和粘贴图像需要一百多毫秒。我正在使用它的程序需要做很多这样的事情，所以加快速度是有益的，而且我还假设这是一种非常标准的操作，所以这段代码是非常可优化的。如果能提供有关如何优化它的任何指导，我将不胜感激。

可能值得一提的是，各种旋转图像通常非常接近，有时会重叠，这就是我进行 mask 的原因，但这是我认为可能效率低下的地方之一。

import cv2
import numpy as np

canvas = np.zeros((2160, 3840, 3), dtype=np.uint8)

img_path = PATH_TO_IMAGE
image = cv2.imread(img_path)

offset_from_center = 10
rotation_angle = 45

width = image.shape[1]
pivot_point = (width/2, offset_from_center)

rotation_mat = cv2.getRotationMatrix2D(pivot_point, -rotation_angle, 1.)

canvas_height = canvas.shape[0]
canvas_width = canvas.shape[1]

rotation_mat[0, 2] += canvas_width/2 - pivot_point[0]
rotation_mat[1, 2] += canvas_height/2 - pivot_point[1]

rotated_image = cv2.warpAffine(image,
                               rotation_mat,
                               (canvas_width, canvas_height))

alpha = np.sum(rotated_image, axis=-1) > 0

alpha = alpha.astype(float)

alpha = np.dstack((alpha, alpha, alpha))

rotated_image = rotated_image.astype(float)
canvas = canvas.astype(float)

foreground = cv2.multiply(alpha, rotated_image)
canvas = cv2.multiply(1.0 - alpha, canvas)

canvas = cv2.add(foreground, canvas)
canvas = canvas.astype(np.uint8)

Answer 1

快速分析显示 np.sum(rotated_image, axis=-1) 特别慢。其他后续操作也有点慢，尤其是数组乘法和 dstack。

---

基于 Numpy 的优化

首先要知道的是，np.sum 会在进行缩减之前自动将数组类型转换为更宽的类型。实际上是一个 64 位 int。这意味着要减少内存中的8 倍大阵列，因此速度会显着降低。话虽这么说，最大的问题是 Numpy 并未针对减少非常小的维度进行优化:它使用非常低效的方法迭代每一行数组，并且迭代比添加 3 个整数要昂贵得多。一种解决方案是手动添加以显着加快速度。这是一个例子:

tmp = rotated_image.astype(np.uint16)
alpha = (tmp[:,:,0] + tmp[:,:,1] + tmp[:,:,2]) > 0

请注意，这是次优的，使用 Cython 或 Numba 可以进一步加快速度(大幅度)。另一种选择是对转置图像进行操作。

然后，当您执行 alpha.astype(float) 时，会使用 64 位 float ，这会占用大量内存，因此任何操作都很慢。这里显然不需要 64 位 float 。可以改用 32 位 float 。事实上，PC GPU 几乎总是使用 32 位 float 来计算图像，因为 64 位运算的成本要高得多(它们需要大量能量和更多晶体管)。

np.dstack 不是必需的，因为 Numpy 可以利用广播(使用 alpha[:,:,None])所以避免大的临时数组 像这样。不幸的是，Numpy 广播在实践中会减慢速度......

cv2.multiply 可以替换为 np.multiply，它有一个特殊的参数 out 用于更快的就地操作(在我的机器上快 30%)。 cv2.multiply 也有一个等效的 dst 参数，正如@ChristophRackwitz 所指出的。这两个函数在我的机器上运行得同样快。

这是一个总结所有这些的更快的实现:

tmp = rotated_image.astype(np.uint16)
alpha = (tmp[:,:,0] + tmp[:,:,1] + tmp[:,:,2]) > 0

alpha = alpha.astype(np.float32)

alpha = np.dstack((alpha, alpha, alpha))

rotated_image = rotated_image.astype(np.float32)
canvas = canvas.astype(np.float32)

foreground = np.multiply(alpha, rotated_image)
np.subtract(1.0, alpha, out=alpha)
np.multiply(alpha, canvas, out=canvas)

np.add(foreground, canvas, out=canvas)
canvas = canvas.astype(np.uint8)

因此，优化后的解决方案快了 4 倍。它仍然远未达到快速，但这将 Numpy 推向了它的极限。

---

使用 Numba 加速代码

创建许多大的临时数组并在小维度上操作远非高效。这可以通过在由 Numba 或 Cython 等 (JIT) 编译器优化的基本循环中计算飞行中的像素来解决，以生成快速的 native 代码。这是一个实现:

import numba as nb

@nb.njit('(uint8[:,:,::1], uint8[:,:,::1])', parallel=True)
def compute(img, canvas):
    for i in nb.prange(img.shape[0]):
        for j in range(img.shape[1]):
            ir = np.float32(img[i, j, 0])
            ig = np.float32(img[i, j, 1])
            ib = np.float32(img[i, j, 2])
            cr = np.float32(canvas[i, j, 0])
            cg = np.float32(canvas[i, j, 1])
            cb = np.float32(canvas[i, j, 2])
            alpha = np.float32((ir + ig + ib) > 0)
            inv_alpha = np.float32(1.0) - alpha
            cr = inv_alpha * cr + alpha * ir
            cg = inv_alpha * cg + alpha * ig
            cb = inv_alpha * cb + alpha * ib
            canvas[i, j, 0] = np.uint8(cr)
            canvas[i, j, 1] = np.uint8(cg)
            canvas[i, j, 2] = np.uint8(cb)

compute(rotated_image, canvas)

这是我的 6 核机器上的性能结果(重复 7 次):

Before:            0.427  s     1x
Optimized Numpy:   0.113  s    ~4x
Numba:             0.0023 s  ~186x

正如我们所见，Numba 实现比优化的 Numpy 实现快得多。如果这还不够，您可以在 GPU 上计算，GPU 对此类任务效率更高。它应该更快并且使用更少的能量，因为 GPU 有专门的单元来进行 alpha 混合(和更宽的 SIMD 单元)。一种解决方案是使用仅适用于 Nvidia GPU 的 CUDA。另一种解决方案是为此使用 OpenCL 甚至 OpenGL。不过，这应该比使用 Numba 复杂一些。