如何使用numba加速python函数

Question

我正在尝试使用 numba 加速我对Floyd-Steinberg 抖动算法的实现。 在阅读了初学者指南后，我在代码中添加了@jit装饰器：

def findClosestColour(pixel):
    colors = np.array([[255, 255, 255], [255, 0, 0], [0, 0, 255], [255, 255, 0], [0, 128, 0], [253, 134, 18]])
    distances = np.sum(np.abs(pixel[:, np.newaxis].T - colors), axis=1)
    shortest = np.argmin(distances)
    closest_color = colors[shortest]
    return closest_color

@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
def floydDither(img_array):
    height, width, _ = img_array.shape
    for y in range(0, height-1):
        for x in range(1, width-1):
            old_pixel = img_array[y, x, :]
            new_pixel = findClosestColour(old_pixel)
            img_array[y, x, :] = new_pixel
            quant_error = new_pixel - old_pixel
            img_array[y, x+1, :] =  img_array[y, x+1, :] + quant_error * 7/16
            img_array[y+1, x-1, :] =  img_array[y+1, x-1, :] + quant_error * 3/16
            img_array[y+1, x, :] =  img_array[y+1, x, :] + quant_error * 5/16
            img_array[y+1, x+1, :] =  img_array[y+1, x+1, :] + quant_error * 1/16
    return img_array

但是，我抛出以下错误：

Untyped global name 'findClosestColour': Cannot determine Numba type of <class 'function'>

我想我知道 numba 不知道findClosestColour的类型，但我刚开始使用 numba 并且不知道如何处理错误。

这是我用来测试该功能的代码：

image = cv2.imread('logo.jpeg')
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
im_out = floydDither(img)

这是我使用的测试图像。

Answer 1

首先，不能从 Numba nopython jitted 函数（又名 njit 函数）调用纯 Python 函数。 这是因为 Numba 需要在编译时跟踪类型以生成高效的二进制文件。

此外，Numba 无法编译表达式pixel[:, np.newaxis].T因为np.newaxis似乎尚不受支持（可能是因为np.newaxis是None ）。 您可以使用pixel.reshape(3, -1).T代替。

请注意，您应该注意类型，因为当两个变量都是np.uint8类型时执行a - b np.uint8导致可能的溢出（例如0 - 1 == 255 ，或者更令人惊讶的是： 0 - 256 = 65280当b是文字整数和a类型的np.uint8 ）。 请注意，该数组是就地计算的，并且像素是在之前写入的

---

尽管 Numba 做得很好，但生成的代码效率不会很高。 您可以使用循环自己迭代颜色以找到最小索引。 这要好一些，因为它不会生成很多小的临时数组。 您还可以指定类型，以便 Numba 提前编译函数。 话虽如此。 这也使代码级别较低，因此更冗长/更难以维护。

这是一个优化的实现：

@nb.njit('int32[::1](uint8[::1])')
def nb_findClosestColour(pixel):
    colors = np.array([[255, 255, 255], [255, 0, 0], [0, 0, 255], [255, 255, 0], [0, 128, 0], [253, 134, 18]], dtype=np.int32)
    r,g,b = pixel.astype(np.int32)
    r2,g2,b2 = colors[0]
    minDistance = np.abs(r-r2) + np.abs(g-g2) + np.abs(b-b2)
    shortest = 0
    for i in range(1, colors.shape[0]):
        r2,g2,b2 = colors[i]
        distance = np.abs(r-r2) + np.abs(g-g2) + np.abs(b-b2)
        if distance < minDistance:
            minDistance = distance
            shortest = i
    return colors[shortest]

@nb.njit('uint8[:,:,::1](uint8[:,:,::1])')
def nb_floydDither(img_array):
    assert(img_array.shape[2] == 3)
    height, width, _ = img_array.shape
    for y in range(0, height-1):
        for x in range(1, width-1):
            old_pixel = img_array[y, x, :]
            new_pixel = nb_findClosestColour(old_pixel)
            img_array[y, x, :] = new_pixel
            quant_error = new_pixel - old_pixel
            img_array[y, x+1, :] =  img_array[y, x+1, :] + quant_error * 7/16
            img_array[y+1, x-1, :] =  img_array[y+1, x-1, :] + quant_error * 3/16
            img_array[y+1, x, :] =  img_array[y+1, x, :] + quant_error * 5/16
            img_array[y+1, x+1, :] =  img_array[y+1, x+1, :] + quant_error * 1/16
    return img_array

naive 版本快 14 倍，而最后一个版本快19 倍。