| """Geometric augmentations for ARC grids.""" |
| import random |
| import torch |
| import numpy as np |
|
|
| class GridAugmentations: |
| def __init__(self, grid_size=64): |
| self.grid_size = grid_size |
| |
| def rotate90(self, grid, k=1): |
| if grid.dim() == 2: |
| return torch.rot90(grid, k=k, dims=(0, 1)) |
| return torch.rot90(grid, k=k, dims=(-2, -1)) |
| |
| def flip_horizontal(self, grid): |
| if grid.dim() == 2: |
| return torch.flip(grid, dims=(0,)) |
| return torch.flip(grid, dims=(-2,)) |
| |
| def flip_vertical(self, grid): |
| if grid.dim() == 2: |
| return torch.flip(grid, dims=(1,)) |
| return torch.flip(grid, dims=(-1,)) |
| |
| def transpose(self, grid): |
| if grid.dim() == 2: |
| return grid.t() |
| return grid.transpose(-2, -1) |
| |
| def random_color_permute(self, grid, colors=None): |
| if colors is None: |
| unique = torch.unique(grid) |
| colors = unique[unique > 0].cpu().numpy().tolist() |
| colors = [c for c in colors if c > 0] |
| if len(colors) <= 1: |
| return grid, {c: c for c in colors} |
| shuffled = colors.copy() |
| random.shuffle(shuffled) |
| cmap = {c: shuffled[i] for i, c in enumerate(colors)} |
| cmap[0] = 0 |
| perm = grid.clone() |
| for old, new in cmap.items(): |
| perm = torch.where(grid == old, torch.tensor(new, dtype=grid.dtype, device=grid.device), perm) |
| return perm, cmap |
| |
| def apply_all_transforms(self, grid, include_color=True): |
| transforms = [(grid, "original")] |
| for k in range(1, 4): |
| transforms.append((self.rotate90(grid, k=k), f"rotate_{k*90}")) |
| transforms.append((self.flip_horizontal(grid), "flip_h")) |
| transforms.append((self.flip_vertical(grid), "flip_v")) |
| transforms.append((self.transpose(grid), "transpose")) |
| if include_color: |
| perm, _ = self.random_color_permute(grid) |
| transforms.append((perm, "color_permute")) |
| return transforms |
|
|
|
|
| def augment_arc_pair(input_grid, output_grid, aug_prob=0.5): |
| if random.random() > aug_prob: |
| return input_grid, output_grid, "none" |
| aug_name = random.choice(["rotate90", "flip_h", "flip_v", "transpose", "color_permute"]) |
| if aug_name == "rotate90": |
| k = random.randint(1, 3) |
| return np.rot90(input_grid, k=k), np.rot90(output_grid, k=k), f"rotate90_{k}" |
| elif aug_name == "flip_h": |
| return np.flip(input_grid, 0), np.flip(output_grid, 0), "flip_h" |
| elif aug_name == "flip_v": |
| return np.flip(input_grid, 1), np.flip(output_grid, 1), "flip_v" |
| elif aug_name == "transpose": |
| return input_grid.T, output_grid.T, "transpose" |
| elif aug_name == "color_permute": |
| colors = np.unique(np.concatenate([input_grid.flatten(), output_grid.flatten()])) |
| colors = colors[colors > 0].tolist() |
| if len(colors) > 1: |
| shuffled = colors.copy() |
| random.shuffle(shuffled) |
| cmap = {c: shuffled[i] for i, c in enumerate(colors)} |
| cmap[0] = 0 |
| return np.vectorize(cmap.get)(input_grid), np.vectorize(cmap.get)(output_grid), "color_permute" |
| return input_grid, output_grid, "none" |
|
|
