perf(color): make cluster cleanup linear and memory-bounded

README.md

@@ -144,6 +144,9 @@ To run tests against the local Moto service, run:
 docker compose --profile testing run --build --rm test
 ```
 
+Tests are split into `test_*.py` files in `vectorizing/tests`. Compose and CI
+run them all. `color_reference.py` keeps the old code for before/after checks.
+
 When running pytest outside Compose, configure the AWS variables for an S3-compatible service first.
 
 ### Required setup

docker-compose.yml

@@ -57,4 +57,4 @@ services:
     depends_on:
       moto-init:
         condition: service_completed_successfully
-    command: python -m pytest vectorizing/tests
+    command: python -m pytest

pytest.ini

@@ -1,3 +1,3 @@
 [pytest]
 pythonpath = .
-python_files = test.py test_*.py *_test.py
+testpaths = vectorizing/tests

vectorizing/solvers/color/quantize.py

@@ -20,12 +20,6 @@ def bilateral_filter(
     return cv2.bilateralFilter(img_arr, d, s, s)
 
 
-def dilate(matrix: np.ndarray) -> np.ndarray:
-    """Dilate labels with a two-by-two kernel and OpenCV's default anchor."""
-    kernel = np.ones((2, 2))
-    return cv2.dilate(matrix, kernel)
-
-
 def fill_holes(matrix: np.ndarray) -> np.ndarray:
     """Replace zero cells with their nearest nonzero value via a distance transform."""
     closest = ndi.distance_transform_edt(
@@ -42,89 +36,50 @@ def enhance(
     labels: np.ndarray,
     colors: np.ndarray,
 ) -> tuple[np.ndarray, np.ndarray]:
-    """Remove components mostly covered by neighboring colors, then fill holes."""
-    dims = img_arr.shape[:2]
-
-    # 0 will be reserved for "unwritten" cells, or holes
+    """Remove boundary-dominated components with O(K*N) work and O(N) image storage."""
+    # Reserve zero for holes, preserving the existing input-dtype shift behavior.
     labels = labels + 1
-
-    clusters = [
-        # Even though we don't allow it, technically the max
-        # value in labels (before adding 1) is 255, so it's safer
-        # to convert to uint16 to avoid overflows
-        np.where(labels == idx + 1, idx + 1, 0).astype(np.uint16)
-        for idx, _ in enumerate(colors)
-    ]
-
-    # For each cluster, store a matrix of its connected components
-    # skimage's default connectivity includes diagonal neighbors.
-    original_connected_components_list = [label(cluster) + 1 for cluster in clusters]
-
-    # Bincounts for each connected components matrix
-    # Meaning bincount[i] = num_appearances(matrix, i)
-    # It's used later to determine what should be erased
-    connected_components_bincounts = [
-        np.bincount(connected_components.flatten())
-        for connected_components in original_connected_components_list
-    ]
-
-    # Make a copy of the connected components
-    # These will be modified and compared with the originals later
-    connected_components_list = [
-        np.array(item, copy=True) for item in original_connected_components_list
-    ]
-
-    for x, cluster_x in enumerate(clusters):
-        dilated_cluster_x = dilate(cluster_x)
-
-        for y, cluster_y in enumerate(clusters):
-            if x == y:
-                continue
-
-            # The overlap between dilated_cluster_x and cluster_y,
-            # meaning, the part of cluster_y that dilated_cluster_x covers
-            overlap = np.logical_and(dilated_cluster_x, cluster_y)
-
-            # Remove overlap from connected components
-            connected_components_list[y] = np.where(
-                overlap,
-                0,
-                connected_components_list[y],
-            )
-
-    for x, connected_components in enumerate(connected_components_list):
-        original_bincount = connected_components_bincounts[x]
-        count = original_bincount.shape[0]
-
-        new_bincount = np.bincount(
-            connected_components.flatten(),
-            minlength=count,  # Force same length
+    valid = (labels > 0) & (labels <= len(colors))
+
+    # OpenCV's 2x2 dilation (default anchor/border) reaches a pixel from above,
+    # left, and upper-left. Mark overlap from any OTHER palette color once,
+    # without treating missing/out-of-palette neighbors as a dilated cluster.
+    boundary = np.zeros(labels.shape, dtype=bool)
+    boundary[1:, :] |= valid[:-1, :] & (labels[1:, :] != labels[:-1, :])
+    boundary[:, 1:] |= valid[:, :-1] & (labels[:, 1:] != labels[:, :-1])
+    boundary[1:, 1:] |= valid[:-1, :-1] & (labels[1:, 1:] != labels[:-1, :-1])
+
+    # The boundary is fixed; reuse its complement for every color's area count.
+    interior = ~boundary
+    del boundary, valid
+    cleaned = np.zeros(img_arr.shape[:2], dtype=np.uint16 if len(colors) else np.uint8)
+    # One scan finds every color's box, including all disconnected components.
+    # max_label excludes invalid labels; zero would instead infer the maximum.
+    regions = ndi.find_objects(labels, max_label=len(colors)) if len(colors) else []
+    for index, region in enumerate(regions):
+        if region is None:
+            continue
+        cluster = labels[region] == index + 1
+        # Keep full (8-neighbor) connectivity, including diagonal contacts.
+        components = label(cluster, connectivity=2)
+        original_counts = np.bincount(components.ravel())
+        interior_counts = np.bincount(
+            # Slice the global mask: other colors outside the box still influence it.
+            components[interior[region]],
+            minlength=len(original_counts),
         )
-
         areas_ratio = np.divide(
-            new_bincount.astype(np.float32),
-            original_bincount.astype(np.float32),
-            out=np.ones((count,), dtype=np.float32),
-            where=original_bincount != 0,
-        )
-
-        # Update cluster
-        # Connected components that lost more than 0.9 of their original area
-        # are completely removed, others are kept
-        clusters[x] = np.where(
-            areas_ratio[original_connected_components_list[x]] <= 0.1,
-            0,
-            clusters[x],
+            interior_counts.astype(np.float32),
+            original_counts.astype(np.float32),
+            out=np.ones(len(original_counts), dtype=np.float32),
+            where=original_counts != 0,
         )
+        # Retain whole components, using the same float32, inclusive <=0.1 cutoff.
+        cleaned[region][cluster & (areas_ratio[components] > 0.1)] = index + 1
+        # Release this map before allocating the next color's component map.
+        del components
 
-    # Update labels with new cluster data
-    labels = np.zeros(dims).astype(np.uint8)
-    for cluster in clusters:
-        labels = np.where(labels == 0, cluster, labels)
-
-    # Fill holes
-    labels = fill_holes(labels)
-    return labels - 1, colors
+    return fill_holes(cleaned) - 1, colors
 
 
 # Try to get the cluster of pixels that represent a transparent background

vectorizing/tests/color_reference.py

@@ -0,0 +1,105 @@
+"""Frozen pre-optimization algorithms used as independent regression oracles.
+
+Keep these intentionally slow implementations separate from production helpers:
+sharing optimized code would let a regression affect both sides of a comparison.
+Cleanup shares only the unchanged fill_holes step so tests can bypass it in both
+paths to inspect component removal before nearest-neighbor filling hides holes.
+"""
+
+import cv2
+import numpy as np
+from PIL import Image
+from skimage.measure import label
+
+from vectorizing.solvers.color import quantize as color_quantize
+
+
+def legacy_initial_centroids(img_arr: np.ndarray, color_count: int) -> np.ndarray:
+    """Preserve the RGB-pixel centroid algorithm from before PR 71."""
+    img = (
+        Image.fromarray(img_arr)
+        .quantize(color_count, method=Image.Quantize.FASTOCTREE)
+        .convert("RGB")
+    )
+    pixels = np.asarray(img)
+    return np.unique(pixels.reshape(-1, pixels.shape[-1]), axis=0).astype(np.uint8)
+
+
+def legacy_create_bitmaps(
+    labels: np.ndarray,
+    colors: np.ndarray,
+    has_background: bool,
+) -> tuple[list[np.ndarray], list[np.ndarray]]:
+    """Preserve the pairwise bitmap accumulation from before PR 73."""
+    bitmaps = [
+        np.where(labels == index, 1, 0).astype(np.uint32)
+        for index in range(len(colors))
+    ]
+    if has_background:
+        bitmaps = bitmaps[1:]
+        colors = colors[1:]
+    zipped = list(zip(bitmaps, colors))
+    zipped = [[bitmap, color] for bitmap, color in zipped if np.sum(bitmap) > 0]
+    bitmaps = [bitmap for bitmap, _ in zipped]
+    colors = [color for _, color in zipped]
+    for x in range(len(bitmaps)):
+        bitmap_x = bitmaps[x]
+        for y in range(x + 1, len(bitmaps)):
+            bitmap_x += bitmaps[y]
+        bitmaps[x] = bitmap_x
+    return bitmaps, colors
+
+
+def legacy_enhance(
+    img_arr: np.ndarray,
+    labels: np.ndarray,
+    colors: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Preserve the pairwise cleanup from main at f69c69b, before PR 74."""
+    dims = img_arr.shape[:2]
+    labels = labels + 1
+    clusters = [
+        np.where(labels == idx + 1, idx + 1, 0).astype(np.uint16)
+        for idx, _ in enumerate(colors)
+    ]
+    # Preserve integer labeling, full connectivity and shifted component IDs.
+    # The optimized path must match this, not redefine the expected behavior.
+    original_connected_components_list = [label(cluster) + 1 for cluster in clusters]
+    connected_components_bincounts = [
+        np.bincount(connected_components.flatten())
+        for connected_components in original_connected_components_list
+    ]
+    connected_components_list = [
+        np.array(item, copy=True) for item in original_connected_components_list
+    ]
+    # Keep the expensive pairwise dilation scans: independence is the point.
+    for x, cluster_x in enumerate(clusters):
+        dilated_cluster_x = cv2.dilate(cluster_x, np.ones((2, 2)))
+        for y, cluster_y in enumerate(clusters):
+            if x == y:
+                continue
+            overlap = np.logical_and(dilated_cluster_x, cluster_y)
+            connected_components_list[y] = np.where(
+                overlap,
+                0,
+                connected_components_list[y],
+            )
+    for x, connected_components in enumerate(connected_components_list):
+        original_bincount = connected_components_bincounts[x]
+        count = original_bincount.shape[0]
+        new_bincount = np.bincount(connected_components.flatten(), minlength=count)
+        areas_ratio = np.divide(
+            new_bincount.astype(np.float32),
+            original_bincount.astype(np.float32),
+            out=np.ones((count,), dtype=np.float32),
+            where=original_bincount != 0,
+        )
+        clusters[x] = np.where(
+            areas_ratio[original_connected_components_list[x]] <= 0.1,
+            0,
+            clusters[x],
+        )
+    labels = np.zeros(dims).astype(np.uint8)
+    for cluster in clusters:
+        labels = np.where(labels == 0, cluster, labels)
+    return color_quantize.fill_holes(labels) - 1, colors