python - 如何使用 numpy 对 Floyd-Steinberg 的抖动算法进行矢量化？

Question

我使用了 Floyd-Steinberg 的抖动算法将图像减少为 6 种颜色的组合。来自 wikipedia page 上的伪代码，我已经能够使用 NumPy 实现该算法。

我很想加快算法速度，并且希望在矢量化我的实现方面得到帮助。这是我最初的实现:

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

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

现在这是我对算法的某些部分进行矢量化的尝试:

def closestColour(img_array):
    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(img_array[..., np.newaxis] - colors.T), axis=2)
    nearest_colors = np.argmin(distances, axis=2)
    quantized = colors[nearest_colors]
    return quantized

def floydDitherVec(img_array):
    new_img = closestColour(img_array)
    error = img_array - new_img
    height, width, _ = new_img.shape

    for y in range(0, height-1):
        for x in range(1, width-1):
            new_img[x+1, y, :] = new_img[x+1, y, :] + (error[x+1, y, :] * 7/16)
            new_img[x-1, y+1, :] = new_img[x-1, y+1, :] + (error[x-1, y+1, :] * 3/16)
            new_img[x, y+1, :] = new_img[x, y+1, :] + (error[x, y+1, :] * 5/16)
            new_img[x+1, y+1, :] = new_img[x+1, y+1, :] + (error[x+1, y+1, :] * 1/16)

    return new_img

最后，我使用 openCV 来读取图像并像这样处理它们:

image = cv2.imread('test2.png')
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

我的总体目标是让算法运行得更快，欢迎任何实现此目标的方法。

Answer 1

我已经能够使用numba来加速算法。实现如下:

@jit(nopython=True)
def floydDitherspeed(img_array):
    height, width, _ = img_array.shape
    colors = np.array([[255, 255, 255], [255, 0, 0], [0, 0, 255], [255, 255, 0], [0, 128, 0], [253, 134, 18]])
    for y in range(0, height-1):
        for x in range(1, width-1):
            old_pixel = img_array[y, x, :]
            max_distance = 195075
            for color in colors:
                p_distances = old_pixel - color
                p_distances = np.power(p_distances, 2)
                distance = p_distances[0] + p_distances[1] + p_distances[2]
                if distance <= max_distance:
                    max_distance = distance
                    new_pixel = color
            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

与 floydDither 函数相比，numba 实现的运行速度大约快 5 倍。