From d0db7e72c399fb23edd1a4bb6e687529defed2f3 Mon Sep 17 00:00:00 2001 From: emerzn <6620417+emerzn@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:19:21 -0400 Subject: [PATCH] fix: minimize total color distance in build_unique_palette Fixes part of inkstitch/inkstitch#2668 which references several problems. The main issues I resolved is inkstitch/inkstitch#1286 and inkstitch/inkstitch#1761. I'm positive it does NOT resolve inkstitch/inkstitch#2612 which I may try tackling in a future PR. **Problem**: when users exported to a .pes file, colors appeared to be swapped. **Root cause**: when converting the user's chosen threads to brother-compatible colors, `build_unique_palette` was greedily assigning palette colors which meant colors later in the threadlist could end up with a worse color than it should have. For example, let's say you want to stitch two threads in this order: `cream --> white` and your machine only has these threads available: `{black,snow}`. Since cream was evaluated first, it would be assigned to snow since that was the closest match, then white would be assigned to black since snow was already taken, resulting in `[snow,black]` when the better solution would have been `[black, snow]` since that pairing minimizes the total color distance across both threads. This is a well-known combinatorics problem: https://en.wikipedia.org/wiki/Assignment_problem Specifically it is of the "unbalanced" flavor since the number of available threads != the number of threads the user is using. **Solution**: use the hungarian algorithm to minimize the combined color shift. This is a rather nasty-looking function but I've purposely kept it generic to help show it is just a standard algorithm. An alternative is to use scipy.optimize.linear_sum_assignment (which itself is just implementing hungarian as well), but pystitch currently has no dependencies and I did not want to introduce scipy just for that. You might think another solution is to simply test every combination, but that would be O(n!) and even with our low numbers like n=64, that is catastrophic! I was trying to avoid adding complexity if there was a simpler solution, but I can't think of any other options. I did corner off the generic algorithm in its own file with some links to explanations, and I commentated on the updated `build_unique_palette` to help explain how it works. --- src/pystitch/Assignment.py | 63 ++++++++++++++++++++++++++++ src/pystitch/EmbThread.py | 73 +++++++++++++++++++++++++-------- test/test_palette.py | 23 ++++++++++- test/test_pec_color_matching.py | 44 ++++++++++++++++++++ 4 files changed, 185 insertions(+), 18 deletions(-) create mode 100644 src/pystitch/Assignment.py create mode 100644 test/test_pec_color_matching.py diff --git a/src/pystitch/Assignment.py b/src/pystitch/Assignment.py new file mode 100644 index 0000000..e3422d6 --- /dev/null +++ b/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 diff --git a/src/pystitch/EmbThread.py b/src/pystitch/EmbThread.py index 0db0ac1..baa25fa 100644 --- a/src/pystitch/EmbThread.py +++ b/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 diff --git a/test/test_palette.py b/test/test_palette.py index b99bd5a..2f46f1b 100644 --- a/test/test_palette.py +++ b/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() diff --git a/test/test_pec_color_matching.py b/test/test_pec_color_matching.py new file mode 100644 index 0000000..8661141 --- /dev/null +++ b/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()