如何使用Numba加速Python实现的Floyd-Steinberg抖动算法并处理报错

问题根因 #

该报错是因为Numba的nopython模式要求所有被调用的函数也必须经过Numba的JIT编译，你定义的findClosestColour函数没有加@jit装饰器，所以被floydDither调用时Numba无法识别该函数的类型，就抛出了这个错误。

解决方法 #

给findClosestColour也加上@jit(nopython=True)装饰器即可，这个函数里用到的numpy操作全部都在Numba的nopython模式支持范围内，不需要额外修改逻辑。同时可以把固定的颜色表提到函数外部，避免每次调用重复创建数组，提升运行效率。

修改后的完整代码如下：

import numpy as np
from numba import jit
import cv2

# 全局颜色常量，避免重复创建
COLORS = np.array([[255, 255, 255], [255, 0, 0], [0, 0, 255], [255, 255, 0], [0, 128, 0], [253, 134, 18]])

@jit(nopython=True)
def findClosestColour(pixel):
    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) # 等价于@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

测试代码优化 #

OpenCV读入的图片默认是uint8类型，直接处理的话浮点误差会被自动截断，建议先转成浮点型再传入函数，处理完成后再转回uint8即可正常使用：

image = cv2.imread('logo.jpeg')
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 转成浮点型避免整数截断
img_float = img.astype(np.float32)
im_out = floydDither(img_float)
# 截断到0-255合法范围后转回uint8
im_out = np.clip(im_out, 0, 255).astype(np.uint8)

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