|
from skimage.color import rgb2lab, deltaE_cie76 |
|
import numpy as np |
|
import cv2 |
|
from sklearn.cluster import MiniBatchKMeans |
|
|
|
def quantize_global(image, k): |
|
k_means = MiniBatchKMeans(n_clusters=k, compute_labels=False, n_init=1) |
|
k_means.fit(image.reshape(-1, 1)) |
|
labels = k_means.predict(image.reshape(-1, 1)) |
|
return np.uint8(k_means.cluster_centers_[labels]).reshape(image.shape) |
|
|
|
def color_transfer(image, old_colors_hex, new_colors_hex): |
|
|
|
|
|
|
|
|
|
|
|
|
|
image_array = np.asarray(image) |
|
image_quantized = quantize_global(image_array, 16) |
|
image_np = np.array(image_quantized) |
|
image_lab = rgb2lab(image_np) |
|
|
|
|
|
|
|
old_colors_lab = [rgb2lab(np.uint8(np.asarray([[list(int(old_color_hex.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))]]))) for old_color_hex in old_colors_hex] |
|
new_colors_rgb = [np.uint8(list(int(new_color_hex.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))) for new_color_hex in new_colors_hex] |
|
|
|
|
|
recolored_image = image.copy() |
|
|
|
|
|
tolerance = 15.0 |
|
|
|
|
|
for old_color_lab, new_color_rgb in zip(old_colors_lab, new_colors_rgb): |
|
|
|
mask = deltaE_cie76(old_color_lab, image_lab) < tolerance |
|
|
|
recolored_image[mask] = new_color_rgb |
|
|
|
|
|
|
|
|
|
recolored_image_bgr = cv2.cvtColor(recolored_image, cv2.COLOR_RGB2BGR) |
|
|
|
|
|
cv2.imwrite('result.jpg', recolored_image_bgr) |
|
|
|
|
|
|
|
|
|
def recolor(image_np, old_colors, new_colors): |
|
image_np = image_np[:, :, :3] |
|
recolored = color_transfer(image_np, old_colors, new_colors) |
|
return recolored |