Update sketch_helper.py

sketch_helper.py

@@ -6,20 +6,19 @@ from skimage.color import lab2rgb
 from sklearn.cluster import KMeans
 
 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)
+    # Determine the number of bins dynamically
+    unique_colors = np.unique(image.reshape(-1, 3), axis=0)
+    n_bins = int(np.ceil(np.sqrt(unique_colors.shape[0])))
+
+    # Cluster the colors using k-means
+    kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(unique_colors)
     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)
+
+    # Replace each pixel with the closest color
+    dists = np.sum((image.reshape(-1, 1, 3) - colors.reshape(1, -1, 3))**2, axis=2)
+    labels = np.argmin(dists, axis=1)
+    return colors[labels].reshape((image.shape[0], image.shape[1], 3)).astype(np.uint8)
+
     
 def get_high_freq_colors(image):
   im = image.getcolors(maxcolors=1024*1024)