Spaces:
Runtime error
Runtime error
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): | |
# Load the image | |
# image = cv2.imread(image_path) | |
# Convert from BGR to RGB | |
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
# Convert image to LAB color space for more accurate color comparison | |
image_array = np.asarray(image) | |
image_quantized = quantize_global(image_array, 16) | |
image_np = np.array(image_quantized) | |
image_lab = rgb2lab(image_np) | |
# image_lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB) | |
# Convert old and new colors to LAB for comparison | |
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] | |
# Create a copy of the image to recolor | |
recolored_image = image.copy() | |
# Define a tolerance for color matching | |
tolerance = 15.0 # Delta E value | |
# For each old color, find the pixels that match and replace them with the new color | |
for old_color_lab, new_color_rgb in zip(old_colors_lab, new_colors_rgb): | |
# Create a mask for pixels matching the old color within the tolerance | |
mask = deltaE_cie76(old_color_lab, image_lab) < tolerance | |
# Replace the colors in the image | |
recolored_image[mask] = new_color_rgb | |
# return recolored_image | |
# Convert the recolored image back to BGR color space | |
recolored_image_bgr = cv2.cvtColor(recolored_image, cv2.COLOR_RGB2BGR) | |
# Save the recolored image | |
# recolored_image_path = '/content/recolored_image.png' | |
cv2.imwrite('result.jpg', recolored_image_bgr) | |
# return recolored_image_path | |
def recolor(image_np, old_colors, new_colors): | |
image_np = image_np[:, :, :3] | |
recolored = color_transfer(image_np, old_colors, new_colors) | |
return recolored |