Update sketch_helper.py

sketch_helper.py

@@ -6,20 +6,18 @@ from skimage.color import lab2rgb
 from sklearn.cluster import KMeans
 
 def color_quantization(image, n_colors):
-    # 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])))
+    # Reshape the image into a 2D array of pixels
+    pixels = image.reshape(-1, 3)
 
-    # Cluster the colors using k-means
-    kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(unique_colors)
-    colors = kmeans.cluster_centers_
+    # Fit k-means clustering algorithm to the pixel data
+    kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(pixels)
 
-    # Replace each pixel with the closest color
-    dists = np.sqrt(np.sum((image.reshape(-1, 1, 3) - colors.reshape(1, -1, 3))**2, axis=2))
-    labels = np.argmin(dists, axis=1)
-    new_image = colors[labels].reshape(image.shape).astype(np.uint8)
-    
-    return new_image
+    # Replace each pixel with its nearest cluster center
+    labels = kmeans.predict(pixels)
+    colors = kmeans.cluster_centers_.astype(np.uint8)
+    quantized_image = colors[labels].reshape(image.shape)
+
+    return quantized_image
 
     
 def get_high_freq_colors(image):