如何基于kmeans聚类加速图像颜色量化并保障色彩还原效果

需求背景 #

• 实现目标：基于K-means聚类完成图像颜色量化，减少图像颜色数量的同时保持视觉效果美观
• 现有问题：当前实现处理1000x1000分辨率、k=32的图像时运行速度极慢，但色彩/色调匹配效果符合业务要求，无等价可替代方案
• 预研优化思路：聚类前先将图像缩小（比如300x300），再将聚类得到的精简调色板应用到原始图像上，暂不具备对应代码实现能力

原始慢版本代码 #

def color_quantize(image, K):
    (h, w) = image.shape[:2]

    img = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    Z = img.reshape((-1, 3))

    # *** my attempt to resize to a smaller size for clustering
    #thumbnail = cv2.resize(img, (300, 300), cv2.INTER_CUBIC)             
    #​Z = thumbnail.reshape((-1, 3))

    Z = np.float32(Z)

    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 15, 1.0)
    ret, label, center = cv2.kmeans(Z, K, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    
    center = np.uint8(center)
    res = center[label.flatten()]

    # *** here should go some backward process to apply the reduced color palette to the original image...

    quantized_img = res.reshape(img.shape)

    quantized_img = cv2.cvtColor(quantized_img, cv2.COLOR_LAB2BGR)
    return quantized_img

最终优化实现方案 #

经过多次测试，最终找到了兼顾着色效果和运行速度的实现：

def color_quantize_fast(image, K):
    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    im_pil = Image.fromarray(np.uint8(img))
    im_pil = im_pil.quantize(K, None, 0, None)
    return cv2.cvtColor(np.array(im_pil.convert("RGB")), cv2.COLOR_RGB2BGR) 

函数说明 #

• 输入：cv2.imread()返回格式的图像
• 输出：同输入格式的量化后图像
• 参数K：目标颜色数量，取值范围为1 <= K <= 255

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