Add object layer: connected components, color splitting, list reducers, overlay/paint/underpaint

itt_solver/object_layer.py

@@ -0,0 +1,309 @@
+"""
+Object extraction and manipulation primitives for ARC-AGI tasks.
+
+Provides connected-component extraction, color-based splitting,
+list reduction (largest/smallest/most_common), spatial queries,
+and composition operations (overlay/paint/underpaint).
+"""
+import numpy as np
+from collections import Counter, deque
+
+
+# ---------------------------------------------------------------------------
+#  Connected component extraction
+# ---------------------------------------------------------------------------
+
+def _flood_fill(grid, start, visited, connectivity=4, univalued=True):
+    """BFS flood fill from start. Returns set of (color, (r, c)) cells."""
+    h, w = grid.shape
+    r0, c0 = start
+    seed_color = int(grid[r0, c0])
+    comp = set()
+    queue = deque([(r0, c0)])
+    visited[r0, c0] = True
+
+    if connectivity == 8:
+        deltas = [(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]
+    else:
+        deltas = [(-1,0),(1,0),(0,-1),(0,1)]
+
+    while queue:
+        r, c = queue.popleft()
+        val = int(grid[r, c])
+        if univalued and val != seed_color:
+            continue
+        comp.add((val, (r, c)))
+        for dr, dc in deltas:
+            nr, nc = r + dr, c + dc
+            if 0 <= nr < h and 0 <= nc < w and not visited[nr, nc]:
+                nval = int(grid[nr, nc])
+                if univalued:
+                    if nval == seed_color:
+                        visited[nr, nc] = True
+                        queue.append((nr, nc))
+                else:
+                    visited[nr, nc] = True
+                    queue.append((nr, nc))
+    return comp
+
+
+def extract_objects(grid, univalued=True, connectivity=4, without_bg=True):
+    """Extract connected components from grid.
+
+    Args:
+        grid: 2D numpy array (int)
+        univalued: if True, each component is single-color
+        connectivity: 4 or 8
+        without_bg: if True, skip the most common color (background)
+
+    Returns:
+        list of objects, each object is a set of (color, (row, col))
+        sorted by size descending
+    """
+    grid = np.array(grid, dtype=int)
+    h, w = grid.shape
+    bg = most_common_color(grid) if without_bg else -1
+    visited = np.zeros((h, w), dtype=bool)
+    objects = []
+
+    for r in range(h):
+        for c in range(w):
+            if visited[r, c]:
+                continue
+            val = int(grid[r, c])
+            if val == bg:
+                visited[r, c] = True
+                continue
+            comp = _flood_fill(grid, (r, c), visited, connectivity, univalued)
+            if comp:
+                objects.append(comp)
+
+    objects.sort(key=len, reverse=True)
+    return objects
+
+
+def split_by_color(grid, without_bg=True):
+    """Split grid into per-color masks. Returns list of (color, grid) pairs
+    where each grid has only that color's pixels (rest = 0)."""
+    grid = np.array(grid, dtype=int)
+    bg = most_common_color(grid) if without_bg else -1
+    colors = sorted(set(grid.flatten()) - {bg})
+    result = []
+    for c in colors:
+        mask_grid = np.zeros_like(grid)
+        mask_grid[grid == c] = c
+        result.append((c, mask_grid))
+    return result
+
+
+# ---------------------------------------------------------------------------
+#  Object to grid conversion
+# ---------------------------------------------------------------------------
+
+def object_to_grid(obj, shape, bg=0):
+    """Render an object (set of (color, (r,c))) onto a grid of given shape."""
+    grid = np.full(shape, bg, dtype=int)
+    for color, (r, c) in obj:
+        if 0 <= r < shape[0] and 0 <= c < shape[1]:
+            grid[r, c] = color
+    return grid
+
+
+def object_to_cropped_grid(obj, bg=0):
+    """Render object cropped to its bounding box."""
+    if not obj:
+        return np.array([[bg]], dtype=int)
+    rows = [r for _, (r, c) in obj]
+    cols = [c for _, (r, c) in obj]
+    rmin, rmax = min(rows), max(rows)
+    cmin, cmax = min(cols), max(cols)
+    h, w = rmax - rmin + 1, cmax - cmin + 1
+    grid = np.full((h, w), bg, dtype=int)
+    for color, (r, c) in obj:
+        grid[r - rmin, c - cmin] = color
+    return grid
+
+
+def normalize_object(obj):
+    """Shift object so its top-left corner is at (0, 0)."""
+    if not obj:
+        return obj
+    rows = [r for _, (r, c) in obj]
+    cols = [c for _, (r, c) in obj]
+    rmin, cmin = min(rows), min(cols)
+    return {(color, (r - rmin, c - cmin)) for color, (r, c) in obj}
+
+
+def shift_object(obj, dr, dc):
+    """Shift all cells by (dr, dc)."""
+    return {(color, (r + dr, c + dc)) for color, (r, c) in obj}
+
+
+# ---------------------------------------------------------------------------
+#  Object queries
+# ---------------------------------------------------------------------------
+
+def object_color(obj):
+    """Color of a univalued object."""
+    colors = {c for c, _ in obj}
+    if len(colors) == 1:
+        return colors.pop()
+    return max(colors, key=lambda c: sum(1 for cc, _ in obj if cc == c))
+
+
+def object_size(obj):
+    return len(obj)
+
+
+def object_bbox(obj):
+    """Returns (rmin, cmin, rmax, cmax)."""
+    rows = [r for _, (r, c) in obj]
+    cols = [c for _, (r, c) in obj]
+    return min(rows), min(cols), max(rows), max(cols)
+
+
+def object_height(obj):
+    rmin, _, rmax, _ = object_bbox(obj)
+    return rmax - rmin + 1
+
+
+def object_width(obj):
+    _, cmin, _, cmax = object_bbox(obj)
+    return cmax - cmin + 1
+
+
+def object_center(obj):
+    rows = [r for _, (r, c) in obj]
+    cols = [c for _, (r, c) in obj]
+    return (sum(rows) / len(rows), sum(cols) / len(cols))
+
+
+# ---------------------------------------------------------------------------
+#  List reducers
+# ---------------------------------------------------------------------------
+
+def largest_object(objects):
+    """Return the largest object by cell count."""
+    return max(objects, key=len) if objects else None
+
+
+def smallest_object(objects):
+    """Return the smallest object by cell count."""
+    return min(objects, key=len) if objects else None
+
+
+def most_common_object(objects):
+    """Return the object whose normalized shape appears most frequently."""
+    if not objects:
+        return None
+    normed = [frozenset(normalize_object(o)) for o in objects]
+    counter = Counter(normed)
+    most_common_shape = counter.most_common(1)[0][0]
+    for o, n in zip(objects, normed):
+        if n == most_common_shape:
+            return o
+    return objects[0]
+
+
+def unique_object(objects):
+    """If exactly one unique normalized shape exists, return it. Else None."""
+    normed = [frozenset(normalize_object(o)) for o in objects]
+    counter = Counter(normed)
+    uniques = [shape for shape, count in counter.items() if count == 1]
+    if len(uniques) == 1:
+        for o, n in zip(objects, normed):
+            if n == uniques[0]:
+                return o
+    return None
+
+
+def filter_by_color(objects, color):
+    """Keep only objects of the given color."""
+    return [o for o in objects if object_color(o) == color]
+
+
+def filter_by_size(objects, size):
+    """Keep only objects of the given size."""
+    return [o for o in objects if len(o) == size]
+
+
+# ---------------------------------------------------------------------------
+#  Color utilities
+# ---------------------------------------------------------------------------
+
+def most_common_color(grid):
+    """Most frequent color in the grid (= background)."""
+    grid = np.array(grid, dtype=int)
+    counts = Counter(grid.flatten().tolist())
+    return counts.most_common(1)[0][0]
+
+
+def least_common_color(grid):
+    """Least frequent color in the grid."""
+    grid = np.array(grid, dtype=int)
+    counts = Counter(grid.flatten().tolist())
+    return counts.most_common()[-1][0]
+
+
+def palette(grid):
+    """Set of all colors in grid."""
+    return set(np.array(grid, dtype=int).flatten().tolist())
+
+
+def color_normalize(grid):
+    """Remap colors by frequency: most common -> 0, next -> 1, etc."""
+    grid = np.array(grid, dtype=int)
+    counts = Counter(grid.flatten().tolist())
+    ranked = [c for c, _ in counts.most_common()]
+    remap = {c: i for i, c in enumerate(ranked)}
+    return np.vectorize(remap.get)(grid)
+
+
+# ---------------------------------------------------------------------------
+#  Composition / overlay
+# ---------------------------------------------------------------------------
+
+def paint(grid, obj):
+    """Paint object onto grid. Object cells OVERWRITE grid cells."""
+    result = np.array(grid, dtype=int).copy()
+    for color, (r, c) in obj:
+        if 0 <= r < result.shape[0] and 0 <= c < result.shape[1]:
+            result[r, c] = color
+    return result
+
+
+def underpaint(grid, obj):
+    """Paint object onto grid, but ONLY where grid has background color."""
+    result = np.array(grid, dtype=int).copy()
+    bg = most_common_color(result)
+    for color, (r, c) in obj:
+        if 0 <= r < result.shape[0] and 0 <= c < result.shape[1]:
+            if result[r, c] == bg:
+                result[r, c] = color
+    return result
+
+
+def overlay_grids(base, foreground):
+    """Overlay foreground onto base. Foreground non-zero pixels overwrite."""
+    base = np.array(base, dtype=int).copy()
+    fg = np.array(foreground, dtype=int)
+    h = min(base.shape[0], fg.shape[0])
+    w = min(base.shape[1], fg.shape[1])
+    mask = fg[:h, :w] != 0
+    base[:h, :w][mask] = fg[:h, :w][mask]
+    return base
+
+
+def cover(grid, obj):
+    """Erase object from grid (replace with background color)."""
+    result = np.array(grid, dtype=int).copy()
+    bg = most_common_color(result)
+    for _, (r, c) in obj:
+        if 0 <= r < result.shape[0] and 0 <= c < result.shape[1]:
+            result[r, c] = bg
+    return result
+
+
+def canvas(bg_color, shape):
+    """Create a blank grid filled with bg_color."""
+    return np.full(shape, bg_color, dtype=int)