| """ |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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} |
|
|
|
|
| |
| |
| |
|
|
| 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)) |
|
|
|
|
| |
| |
| |
|
|
| 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] |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
