问如何用numba加速python函数
EN

Stack Overflow用户

提问于 2021-10-27 02:53:20

回答 1查看 214关注 0票数 0

我正试图用numba加速抖动算法的实现。在阅读了初学者指南之后，我在代码中添加了@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)

这是我用过的测试图像。

python

numpy

numba

人脸识别特惠活动

精准实时的人脸检测、分析、识别等服务，在线、离线满足多样需求，可应用于智慧零售、智慧社区、在线娱乐等多种应用场景

回答 1

Stack Overflow用户

回答已采纳

发布于 2021-10-27 13:23:30

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

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

请注意，您应该小心这些类型，因为当两个变量都是np.uint8类型时执行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