Update sketch_helper.py

sketch_helper.py

@@ -15,9 +15,11 @@ def color_quantization(image, n_colors):
     colors = kmeans.cluster_centers_
 
     # Replace each pixel with the closest color
-    dists = np.sum((image.reshape(-1, 1, 3) - colors.reshape(1, -1, 3))**2, axis=2)
+    dists = np.sqrt(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)
+    new_image = colors[labels].reshape(image.shape).astype(np.uint8)
+    
+    return new_image)
 
     
 def get_high_freq_colors(image):