Update sketch_helper.py

sketch_helper.py

@@ -1,7 +1,25 @@
 import numpy as np
 import cv2
 from PIL import Image
-
+from skimage.color import rgb2lab
+from skimage.color import lab2rgb
+
+def color_quantization(image, n_colors):
+    # Convert image to LAB color space
+    lab_image = rgb2lab(image)
+    # Reshape image to 2D array of pixels
+    pixels = lab_image.reshape(-1, 3)
+    # Perform K-means clustering
+    kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(pixels)
+    # Replace each pixel with the closest color
+    labels = kmeans.predict(pixels)
+    colors = kmeans.cluster_centers_
+    quantized_pixels = colors[labels]
+    # Convert quantized image back to RGB color space
+    quantized_lab_image = quantized_pixels.reshape(lab_image.shape)
+    quantized_rgb_image = lab2rgb(quantized_lab_image)
+    return (quantized_rgb_image * 255).astype(np.uint8)
+    
 def get_high_freq_colors(image):
   im = image.getcolors(maxcolors=1024*1024)
   sorted_colors = sorted(im, key=lambda x: x[0], reverse=True)
@@ -24,27 +42,6 @@ def color_quantization_old(image, n_colors):
     labels = np.argmin(dists, axis=1)
     return colors[labels].reshape((image.shape[0], image.shape[1], 3)).astype(np.uint8)
 
-def color_quantization(image, n_colors=8, rounds=1):
-    h, w = image.shape[:2]
-    samples = np.zeros([h*w,3], dtype=np.float32)
-    count = 0
-
-    for x in range(h):
-        for y in range(w):
-            samples[count] = image[x][y]
-            count += 1
-
-    compactness, labels, centers = cv2.kmeans(samples,
-            n_colors, 
-            None,
-            (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10000, 0.0001), 
-            rounds, 
-            cv2.KMEANS_RANDOM_CENTERS)
-
-    centers = np.uint8(centers)
-    res = centers[labels.flatten()]
-    return res.reshape((image.shape))
-
 def create_binary_matrix(img_arr, target_color):
     # Create mask of pixels with target color
     mask = np.all(img_arr == target_color, axis=-1)