| """Label-free raw prediction for one or more Visual Sudoku boards.""" |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from omegaconf import DictConfig |
| from softasp_solver.sudoku.softasp_sudoku import SoftASPSudoku |
|
|
| from softasp.instantiations.checkpoint import read_checkpoint, validate_checkpoint |
| from softasp.sudoku_reporting import file_sha256 |
|
|
|
|
| def load_model( |
| checkpoint: Path, device: str = "auto", *, cfg: DictConfig | None = None |
| ) -> tuple[SoftASPSudoku, torch.device]: |
| """Load public inference weights; optimizer state and labels are unnecessary.""" |
| target = torch.device( |
| ("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device |
| ) |
| saved = read_checkpoint(checkpoint) |
| if cfg is not None: |
| validate_checkpoint( |
| cfg, |
| saved, |
| legacy_model="softasp_sudoku", |
| legacy_constraints={"_target_": "softasp_solver.sudoku.constraints.TPOperatorLoss"}, |
| ) |
| model_config = saved["config"]["model"] |
| if model_config["name"] != "softasp_sudoku": |
| raise ValueError("Expected a Sudoku AS2 checkpoint") |
| params = dict(model_config["params"]) |
| if params.get("oracle_perception", False): |
| raise ValueError("Oracle checkpoints are not supported") |
| model = SoftASPSudoku(**params) |
| if cfg is not None: |
| from hydra.utils import instantiate |
|
|
| model.constraint_loss = instantiate(cfg.instantiation.constraints) |
| model.load_state_dict(saved["model_state_dict"], strict=True) |
| return model.to(target).eval(), target |
|
|
|
|
| @torch.no_grad() |
| def predict_file( |
| checkpoint: Path, input_path: Path, device: str = "auto", *, cfg: DictConfig | None = None |
| ) -> dict: |
| """Read images and clue positions only; output raw digits 1 through 9.""" |
| with np.load(input_path, allow_pickle=False) as data: |
| images = data["images"].copy() |
| mask = data["puzzle_mask"].copy() |
| if images.ndim != 4 or images.shape[1:] != (252, 252, 1) or len(images) == 0: |
| raise ValueError("images must have shape (N, 252, 252, 1), N > 0") |
| if images.dtype != np.uint8 and not ( |
| images.dtype == np.float32 |
| and np.isfinite(images).all() |
| and np.all((images >= 0) & (images <= 1)) |
| ): |
| raise ValueError("images must be uint8 [0,255] or float32 [0,1]") |
| if mask.shape != (len(images), 81) or not np.isin(mask, [0, 1]).all(): |
| raise ValueError("puzzle_mask must have shape (N,81) and contain only 0/1") |
| model, target = load_model(checkpoint, device, cfg=cfg) |
| predictions = [] |
| |
| for start in range(0, len(images), 32): |
| out = model.predict( |
| torch.from_numpy(images[start : start + 32]).to(target), |
| torch.from_numpy(mask[start : start + 32].astype(np.int64)).to(target), |
| ) |
| logits = out["logits"] |
| if logits.shape[1:] != (81, 9) or not torch.isfinite(logits).all(): |
| raise ValueError("Expected finite raw Sudoku logits") |
| predictions.append(logits.argmax(-1).cpu()) |
| digits = torch.cat(predictions).numpy().reshape(-1, 9, 9) + 1 |
| expected = np.arange(1, 10) |
| valid = [] |
| for board in digits: |
| boxes = board.reshape(3, 3, 3, 3).transpose(0, 2, 1, 3).reshape(9, 9) |
| valid.append( |
| bool( |
| np.all(np.sort(board, axis=1) == expected) |
| and np.all(np.sort(board.T, axis=1) == expected) |
| and np.all(np.sort(boxes, axis=1) == expected) |
| ) |
| ) |
| return { |
| "decoder": "raw_argmax", |
| "digits": digits.tolist(), |
| "structurally_valid": valid, |
| "checkpoint_sha256": file_sha256(checkpoint), |
| "input_sha256": file_sha256(input_path), |
| } |
|
|
