Update sketch_helper.py

sketch_helper.py

@@ -12,7 +12,7 @@ def get_high_freq_colors(image):
   high_freq_colors = [c for c in sorted_colors if c[0] > max(2, mean_freq)]  # Ignore colors that occur very few times (less than 2) or less than half the average frequency
   return high_freq_colors
 
-def color_quantization(image, n_colors):
+def color_quantization_old(image, n_colors):
     # Get color histogram
     hist, _ = np.histogramdd(image.reshape(-1, 3), bins=(256, 256, 256), range=((0, 256), (0, 256), (0, 256)))
     # Get most frequent colors
@@ -24,6 +24,27 @@ def color_quantization(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,
+            clusters, 
+            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)