reassigning colors

colorbynumber/assign_colors.py

@@ -0,0 +1,27 @@
+from scipy.optimize import linear_sum_assignment
+import numpy as np
+
+def assign_colors(palette_a, palette_b, weights):
+    """Given:
+    - color palette A generated by k-means,
+    - color palette B provided by the user and
+    - weight w for each color in A corresponding to "importance" of the color,
+    assigns each color in B to a color in A such that the sum of the squared
+    distances between the colors in A and B is minimized.
+    Returns a dictionary mapping each color in B to a color in A.
+    """
+    # Create a cost matrix where the element at position (i, j) is the
+    # squared distance between color i in palette A and color j in palette B,
+    # multiplied by the weight of color i.
+    cost_matrix = np.zeros((len(palette_a), len(palette_b)))
+    for i, color_a in enumerate(palette_a):
+        for j, color_b in enumerate(palette_b):
+            cost_matrix[i, j] = weights[i] * np.sum((np.array(color_a) - np.array(color_b))**2)
+
+    # Use the linear_sum_assignment function to find the optimal assignment.
+    row_indices, col_indices = linear_sum_assignment(cost_matrix)
+
+    # Create a dictionary mapping each color in B to a color in A.
+    color_mapping = {tuple(palette_b[j]): tuple(palette_a[i]) for i, j in zip(row_indices, col_indices)}
+
+    return color_mapping

environment.yml

@@ -8,5 +8,6 @@ dependencies:
   - opencv
   - jupyterlab
   - matplotlib
+  - scipy
   - pip:
     - python-polylabel