fix: minimize total color distance in build_unique_palette

src/pystitch/Assignment.py

@@ -0,0 +1,63 @@
+from typing import Sequence
+
+
+# This project doesn't run flake8 yet, but inkstitch's does and this function
+# trips C901 there. The branching is inherent to the shortest-augmenting-path
+# formulation of the algorithm, not something worth splitting up just to
+# satisfy the linter.
+def minimal_assignment(cost: Sequence[Sequence[float]]) -> list[int]:  # noqa: C901
+    """Solves the assignment problem: given cost[row][column] with
+    len(rows) <= len(columns), returns for each row a distinct column such
+    that the total cost is minimal (Hungarian algorithm).
+    See: https://en.wikipedia.org/wiki/Hungarian_algorithm
+
+    Good visualization: https://www.youtube.com/watch?v=cQ5MsiGaDY8
+    but note we are storing an extra matrix (row_potential x column_potential)
+    instead of constantly mutating every value in the cost matrix.
+    """
+    rows = len(cost)
+    columns = len(cost[0])
+    infinity = float("inf")
+    # Potentials and matching, 1-indexed with column 0 as scratch space.
+    row_potential: list[float] = [0] * (rows + 1)
+    column_potential: list[float] = [0] * (columns + 1)
+    match = [0] * (columns + 1)  # match[column] = row assigned to it
+    path = [0] * (columns + 1)
+    for row in range(1, rows + 1):
+        match[0] = row
+        j0 = 0
+        min_value = [infinity] * (columns + 1)
+        used = [False] * (columns + 1)
+        while True:
+            used[j0] = True
+            i0 = match[j0]
+            delta = infinity
+            j1 = 0
+            for j in range(1, columns + 1):
+                if used[j]:
+                    continue
+                current = cost[i0 - 1][j - 1] - row_potential[i0] - column_potential[j]
+                if current < min_value[j]:
+                    min_value[j] = current
+                    path[j] = j0
+                if min_value[j] < delta:
+                    delta = min_value[j]
+                    j1 = j
+            for j in range(0, columns + 1):
+                if used[j]:
+                    row_potential[match[j]] += delta
+                    column_potential[j] -= delta
+                else:
+                    min_value[j] -= delta
+            j0 = j1
+            if match[j0] == 0:
+                break
+        while j0 != 0:
+            j1 = path[j0]
+            match[j0] = match[j1]
+            j0 = j1
+    assignment = [0] * rows
+    for j in range(1, columns + 1):
+        if match[j] != 0:
+            assignment[match[j] - 1] = j - 1
+    return assignment

src/pystitch/EmbThread.py

@@ -1,20 +1,59 @@
-from typing import Optional
-
-def build_unique_palette(thread_palette, threadlist):
-    """Turns a threadlist into a unique index list with the thread palette"""
-    chart = [None] * len(thread_palette)  # Create a lookup chart.
-    for thread in set(
-        threadlist
-    ):  # for each unique color, move closest remaining thread to lookup chart.
-        index = thread.find_nearest_color_index(thread_palette)
-        if index is None:
-            break  # No more threads remain in palette
-        thread_palette[index] = None  # entries may not be reused.
-        chart[index] = thread  # assign the given index to the lookup.
-
-    palette = []
-    for thread in threadlist:  # for each thread, return the index.
-        palette.append(thread.find_nearest_color_index(chart))
+from typing import Optional, Sequence
+
+from .Assignment import minimal_assignment
+
+
+def build_unique_palette(
+    thread_palette: list[Optional["EmbThread"]],
+    threadlist: Sequence["EmbThread"],
+) -> list[Optional[int]]:
+    """Turns a threadlist into a unique index list with the thread palette.
+
+    Each unique thread gets a distinct palette entry, chosen so that the
+    total color distance of the whole assignment is minimal. This is a classic
+    Assignment Problem (https://en.wikipedia.org/wiki/Assignment_problem).
+    If there are more unique threads than palette entries, the extra threads
+    reuse the entry of the nearest assigned thread.
+    """
+    unique_threadlist = []  # removes duplicates
+    for thread in threadlist:
+        if thread not in unique_threadlist:
+            unique_threadlist.append(thread)
+    valid_entries = [
+        (index, entry) for index, entry in enumerate(thread_palette) if entry is not None
+    ]
+    valid_indices = [index for index, _ in valid_entries]
+
+    # create two-way lookups dicts so we cache which of our threads are assigned to which thread_palette index
+    index_to_thread: list[Optional["EmbThread"]] = [None] * len(thread_palette)
+    thread_to_index: dict["EmbThread", int] = {}
+    # if we have more unique desired threads than palette threads, then assignable is shortened
+    assignable = unique_threadlist[: len(valid_entries)]
+    if assignable:
+        # Build a cost matrix for the minimal_assignment algorithm
+        cost = [
+            [
+                color_distance_red_mean(
+                    thread.get_red(), thread.get_green(), thread.get_blue(),
+                    entry.get_red(), entry.get_green(), entry.get_blue(),
+                )
+                for index, entry in valid_entries
+            ]
+            for thread in assignable
+        ]
+        # Use minimal_assignment to map desired threads to threads available in the palette
+        for row, column in enumerate(minimal_assignment(cost)):
+            index = valid_indices[column]
+            index_to_thread[index] = assignable[row]
+            thread_to_index[assignable[row]] = index
+
+    palette: list[Optional[int]] = []
+    for thread in threadlist:  # for each thread, return its assigned index.
+        palette_index = thread_to_index.get(thread)
+        if palette_index is None:
+            # This thread never got its own slot, so reuse the nearest one
+            palette_index = thread.find_nearest_color_index(index_to_thread)
+        palette.append(palette_index)
     return palette
 
 

test/test_palette.py

@@ -1,7 +1,7 @@
 import unittest
 
 from pystitch import EmbPattern, EmbThread
-from pystitch.EmbThreadPec import get_thread_set
+from pystitch.EmbThreadPec import get_thread_set, EmbThreadPec
 from pystitch.EmbThread import build_unique_palette, build_nonrepeat_palette, build_palette
 
 
@@ -19,6 +19,12 @@ def test_unique_palette(self):
         self.assertNotEqual(palette[0], palette[3], "Red and altered Red")
         self.assertEqual(palette[1], palette[2], "Blue and Blue")
 
+    def test_unique_palette_same_thread_repeated_shares_an_index(self):
+        red = EmbThread({"rgb": (255, 0, 0), "name": "Red"})
+        threadset = get_thread_set()
+        palette = build_unique_palette(threadset, [red, EmbThread(red)])
+        self.assertEqual(palette[0], palette[1])
+
     def test_unique_palette_large(self):
         """Excessive palette entries that all map, should be mapped"""
         pattern = EmbPattern()
@@ -52,6 +58,21 @@ def test_unique_palette_max(self):
         for i in range(1, len(palette)):
             self.assertNotEqual(palette[i-1], palette[i])
 
+    def test_unique_palette_order_invariant(self):
+        """thread assignment should not be greedy, it should minimize the total shift of colors.
+        i.e just because eggshell comes first and its closest color is white, it shouldn't be
+        assigned to white, as there is a better match later in the threadlist (actual white)
+        """
+        eggshell = EmbThread({"rgb": (255, 255, 245), "name": "Eggshell", "catalog": "0101"})
+        white = EmbThread({"rgb": (255, 255, 255), "name": "White", "catalog": "0015"})
+        machine_black = EmbThreadPec(0, 0, 0, "Black", "20")
+        machine_white = EmbThreadPec(255, 255, 255, "White", "1")
+        palette = build_unique_palette([machine_black, machine_white], [eggshell, white])
+        self.assertEqual(0, palette[0],
+                            "Eggshell should be assigned to the Black thread")
+        self.assertEqual(1, palette[1],
+                         "White should be assigned to the White thread")
+
     def test_nonrepeat_palette_moving(self):
         """The almost same color should not get plotted to the same palette index"""
         pattern = EmbPattern()

test/test_pec_color_matching.py

@@ -0,0 +1,44 @@
+import io
+import unittest
+
+from pystitch import EmbPattern, EmbThread, write_pec
+from pystitch.PecReader import read_pec
+
+
+def pattern_with_blocks(threads):
+    pattern = EmbPattern()
+    for thread in threads:
+        pattern.add_block([(0, 0), (0, 100), (100, 100)], thread)
+    return pattern
+
+
+def write_then_read_pec(pattern):
+    stream = io.BytesIO()
+    write_pec(pattern, stream)
+    stream.seek(8)  # skip the "#PEC0001" signature; read_pec starts at "LA:"
+    result = EmbPattern()
+    read_pec(stream, result)
+    return result
+
+
+class TestPecColorMatching(unittest.TestCase):
+
+    def test_eggshell_then_white_is_not_swapped(self):
+        eggshell = EmbThread({"rgb": (255, 255, 245), "name": "Eggshell", "catalog": "0101"})
+        white = EmbThread({"rgb": (255, 255, 255), "name": "White", "catalog": "0015"})
+        pattern = write_then_read_pec(pattern_with_blocks([eggshell, white]))
+        self.assertEqual(len(pattern.threadlist), 2)
+        self.assertNotEqual(pattern.threadlist[0].description, "White",  # it's Flesh Pink instead
+                            "the Eggshell block should be shown as an off-white, not White")
+        self.assertEqual(pattern.threadlist[1].description, "White",
+                         "the White block should be shown as White")
+
+    def test_same_thread_repeated_shares_an_index(self):
+        red = EmbThread({"rgb": (255, 0, 0), "name": "Red"})
+        pattern = write_then_read_pec(pattern_with_blocks([red, EmbThread(red)]))
+        self.assertEqual(len(pattern.threadlist), 2)
+        self.assertEqual(pattern.threadlist[0].description, pattern.threadlist[1].description)
+
+
+if __name__ == "__main__":
+    unittest.main()