Spaces:
Running on Zero
Running on Zero
| """Render Sudoku grids as images of MNIST digits and build training/eval batches. | |
| Image: 1 x 144 x 144 in [-1, 1] (9 x 9 cells of 16 px, the paper's training resolution). | |
| Conditioning: the solution image with non-clue cells blacked out, plus a pixel mask of clue cells. | |
| """ | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from . import sudoku | |
| CELL = 16 | |
| IMG = 9 * CELL | |
| class DigitBank: | |
| """MNIST digits resized to CELL x CELL, grouped by class: images[d] is (N_d, CELL, CELL) in [-1, 1].""" | |
| def __init__(self, train=True, data_dir="data", images=None, labels=None): | |
| if images is None: | |
| from torchvision.datasets import MNIST | |
| ds = MNIST(data_dir, train=train, download=True) | |
| images, labels = ds.data, ds.targets | |
| x = F.interpolate(images[:, None].float() / 255.0, size=(CELL, CELL), mode="area")[:, 0] | |
| x = x * 2 - 1 | |
| n = min(int((labels == d).sum()) for d in range(10)) | |
| self.images = torch.stack([x[labels == d][:n] for d in range(10)]) # (10, n, CELL, CELL) | |
| def to(self, device): | |
| self.images = self.images.to(device) | |
| return self | |
| def render(grids, bank, generator=None): | |
| """grids: (B, 9, 9) long tensor of digits 1-9 (0 renders as black). Returns (B, 1, IMG, IMG).""" | |
| B = grids.shape[0] | |
| flat = grids.reshape(B, 81).to(bank.images.device) | |
| idx = torch.randint(bank.images.shape[1], (B, 81), generator=generator, device="cpu").to(flat.device) | |
| cells = bank.images[flat, idx] # (B, 81, CELL, CELL) | |
| cells = torch.where((flat == 0)[..., None, None], torch.full_like(cells, -1.0), cells) | |
| return cells.reshape(B, 9, 9, CELL, CELL).permute(0, 1, 3, 2, 4).reshape(B, 1, IMG, IMG) | |
| def cells(images): | |
| """Inverse of the layout in render: (B, 1, IMG, IMG) -> (B, 81, CELL, CELL).""" | |
| B = images.shape[0] | |
| return images.reshape(B, 9, CELL, 9, CELL).permute(0, 1, 3, 2, 4).reshape(B, 81, CELL, CELL) | |
| def pixel_mask(puzzles): | |
| """(B, 9, 9) puzzles -> (B, 1, IMG, IMG) float mask, 1 on clue cells.""" | |
| m = (puzzles > 0).float() | |
| return m.repeat_interleave(CELL, 1).repeat_interleave(CELL, 2)[:, None] | |
| def make_batch(puzzles, solutions, bank, generator=None): | |
| """Render a batch. Clue pixels in cond are the same digit images as in x0.""" | |
| puzzles, solutions = torch.as_tensor(puzzles).long(), torch.as_tensor(solutions).long() | |
| x0 = render(solutions, bank, generator) | |
| mask = pixel_mask(puzzles).to(x0.device) | |
| cond = torch.where(mask > 0, x0, torch.full_like(x0, -1.0)) | |
| return dict(x0=x0, cond=cond, mask=mask, puzzles=puzzles, solutions=solutions) | |
| class PuzzleSource: | |
| """Samples (puzzles, solutions) for 'extreme' (TRM csv), 'hard' (0-26 random clues of a real solution) | |
| or 'solutions' (no clues; Painter data). Solutions always come from the sudoku-extreme split, so train and | |
| test grids are disjoint; training samples get a random validity-preserving augmentation.""" | |
| def __init__(self, mode, split="train", data_dir="data", seed=0): | |
| if mode not in ("extreme", "hard", "solutions"): | |
| raise ValueError(mode) | |
| self.mode, self.split, self.rng = mode, split, np.random.default_rng(seed) | |
| self.puzzles, self.solutions = sudoku.load_extreme(split, data_dir) | |
| def sample(self, n): | |
| i = self.rng.integers(len(self.solutions), size=n) | |
| p, s = sudoku.augment(self.rng, self.puzzles[i], self.solutions[i]) | |
| if self.mode == "hard": | |
| p = np.stack([sudoku.hard_puzzle(self.rng, g)[0] for g in s]) | |
| elif self.mode == "solutions": | |
| p = np.zeros_like(p) | |
| return p, s | |
| def fixed(self, n, seed=1234): | |
| """A deterministic evaluation set (no augmentation).""" | |
| rng = np.random.default_rng(seed) | |
| i = rng.choice(len(self.solutions), size=n, replace=False) | |
| p, s = self.puzzles[i].astype(np.int64), self.solutions[i].astype(np.int64) | |
| if self.mode == "hard": | |
| p = np.stack([sudoku.hard_puzzle(rng, g)[0] for g in s]) | |
| return p, s | |