如何使用numpy向量化Floyd-Steinberg的抖动算法？

admin

我使用了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)

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

转载请注明出处:http://www.intsu.net/article/20230518/2247543.html