Update sketch_helper.py

sketch_helper.py

@@ -31,17 +31,18 @@ def get_high_freq_colors(image, similarity_threshold=30):
     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)
+    # Convert the color frequency list to a set of unique colors
+    unique_colors = set([color for _, color in color_frequency_list])
+    
+    # Create a mask for the image with True where the color is in the unique colors set
+    mask = np.any(np.all(image.reshape(-1, 1, 3) == np.array(list(unique_colors)), axis=2), axis=1).reshape(image.shape[:2])
+    
+    # Create a new image with all pixels set to white
+    new_image = np.full_like(image, 255)
     
-    return image
+    # Copy the pixels from the original image that have a color in the color frequency list
+    new_image[mask] = image[mask]
+    return new_image
 
 def create_binary_matrix(img_arr, target_color):
     # Create mask of pixels with target color