图像RGB转Lab通道，并且量化ab通道

如题

import cv2
import matplotlib.pyplot as plt
import numpy as np


def quantized_color_preprocess(image, centroids):
    h, w, c = image.shape
    image = cv2.cvtColor(image, cv2.COLOR_BGR2Lab)  # BGR to Lab
    ab = image[:,:,1:]

    a = np.argmin(np.linalg.norm(centroids[None, :, :] - ab.reshape([-1,2])[:, None, :], axis=2),axis=1)
    # 256 256  quantized color (4bit)

    quantized_ab = a.reshape([-1, h, w])
    return quantized_ab

centroids = np.load('centroids_16k_kinetics_10000samples.npy', allow_pickle=True)
image = cv2.imread("../../data/DAVIS/JPEGImages/480p/cat-girl/00030.jpg")
image_lab = cv2.cvtColor(image, cv2.COLOR_RGB2Lab)
image_quan = quantized_color_preprocess(image_lab, centroids)

fig = plt.figure(figsize=(16, 3))
plt.subplot(131)
plt.imshow(image[:,:,::-1])
plt.title("RGB image")
plt.subplot(132)
plt.imshow(image_lab)
plt.title("Lab image")
plt.subplot(133)
plt.imshow(image_quan[0])
plt.title("Quantized ab channels to 16 classes")

结果