Update sketch_helper.py

sketch_helper.py

@@ -6,27 +6,42 @@ from skimage.color import lab2rgb
 from sklearn.cluster import KMeans
 
     
-def get_high_freq_colors(image):
-  im = image.getcolors(maxcolors=1024*1024)
-  sorted_colors = sorted(im, key=lambda x: x[0], reverse=True)
-  
-  freqs = [c[0] for c in sorted_colors]
-  mean_freq = sum(freqs) / len(freqs)
+def count_high_freq_colors(image):
+    im = image.getcolors(maxcolors=1024*1024)
+    sorted_colors = sorted(im, key=lambda x: x[0], reverse=True)
 
-  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
+    freqs = [c[0] for c in sorted_colors]
+    mean_freq = sum(freqs) / len(freqs)
 
-def color_quantization(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
-    colors = np.argwhere(hist > 0)
-    colors = colors[np.argsort(hist[colors[:, 0], colors[:, 1], colors[:, 2]])[::-1]]
-    colors = colors[:n_colors]
-    # 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)
+    high_freq_colors = [c for c in sorted_colors if c[0] > max(2, mean_freq*1.25)]
+    return high_freq_colors
+
+def get_high_freq_colors(image, similarity_threshold=30):
+    image_copy = image.copy()
+    high_freq_colors = count_high_freq_colors(image)
+    # Check for similar colors and replace the lower frequency color with the higher frequency color in the image
+    for i, (freq1, color1) in enumerate(high_freq_colors):
+        for j, (freq2, color2) in enumerate(high_freq_colors):      
+            if (color_distance(color1, color2) < similarity_threshold) or (color_distance(color1, opaque_color_on_white(color2, 0.5)) < 5):
+              if(freq2 > freq1):
+                replace_color(image_copy, color1, color2)
+            
+    high_freq_colors = count_high_freq_colors(image_copy)
+    print(high_freq_colors)
+    return [high_freq_colors, image_copy]
+    
+def color_quantization(image, color_frequency_list):
+    # Extract the color values from the frequency list
+    color_values = [color for _, color in color_frequency_list]
+
+    # Replace the colors that are not in the frequency list with white
+    mask = np.ones(image.shape[:2], dtype=bool)
+    for color in color_values:
+        color_mask = np.all(image == color, axis=2)
+        mask = np.logical_and(mask, np.logical_not(color_mask))
+    image[mask] = (255, 255, 255)
+    
+    return image
 
 def create_binary_matrix(img_arr, target_color):
     # Create mask of pixels with target color