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()