| |
| """Wave 1 static spatial remapping solvers. |
| |
| A4: downsample_stride — strided sampling of input |
| A7: symmetry_complete — mirror to complete L-R or T-B symmetry |
| A1: extract_inner — remove border frame |
| A2: add_border — add constant border |
| A6: sparse_fill — pixel to block expansion |
| B1: channel_filter — keep only certain colors |
| |
| Scan results (2026-04-27): 0 arc-gen validated matches. |
| Kept for future tasks and as building blocks. |
| """ |
|
|
| import numpy as np |
| from ..data_loader import get_exs, fixed_shapes |
| from ..gather_helpers import _build_gather_model, _build_gather_model_with_const |
| from ..onnx_helpers import mk, _make_int64_init, _build_pad_node, add_onehot_block |
| from ..constants import GH, GW |
|
|
|
|
| def s_downsample_stride(td): |
| """out[r,c] = inp[r*sH + oH, c*sW + oW] for integer strides.""" |
| exs = get_exs(td) |
| sp = fixed_shapes(td) |
| if sp is None: |
| return None |
| (IH, IW), (OH, OW) = sp |
| if OH >= IH or OW >= IW: |
| return None |
|
|
| for sh in range(2, 6): |
| for sw in range(2, 6): |
| for oh_off in range(sh): |
| for ow_off in range(sw): |
| ok = True |
| for inp, out in exs: |
| sampled = inp[oh_off::sh, ow_off::sw] |
| if sampled.shape != out.shape or not np.array_equal(sampled, out): |
| ok = False |
| break |
| if ok: |
| idx = np.zeros((OH, OW, 2), dtype=np.int64) |
| for r in range(OH): |
| for c in range(OW): |
| idx[r, c] = [r * sh + oh_off, c * sw + ow_off] |
| return _build_gather_model(OH, OW, idx) |
| return None |
|
|
|
|
| def s_symmetry_complete(td): |
| """Complete partial T-B symmetry by adding mirrored + original via Gather.""" |
| from onnx import helper, numpy_helper |
|
|
| exs = get_exs(td) |
| sp = fixed_shapes(td) |
| if sp is None: |
| return None |
| (IH, IW), (OH, OW) = sp |
| if (IH, IW) != (OH, OW): |
| return None |
| if IH < 2: |
| return None |
|
|
| |
| ok = True |
| for inp, out in exs: |
| exp = inp.copy() |
| for r in range(IH // 2): |
| for c in range(IW): |
| v = max(int(inp[r, c]), int(inp[IH - 1 - r, c])) |
| exp[r, c] = v |
| exp[IH - 1 - r, c] = v |
| if not np.array_equal(out, exp): |
| ok = False |
| break |
|
|
| if ok: |
| |
| pad_h, pad_w = GH - OH, GW - OW |
| mirror_idx = np.zeros((GH * GW,), dtype=np.int64) |
| mask = np.zeros((1, 1, GH, GW), dtype=np.float32) |
| self_idx = np.zeros((GH * GW,), dtype=np.int64) |
| for r in range(OH): |
| for c in range(OW): |
| self_idx[r * GW + c] = r * GW + c |
| mirror_idx[r * GW + c] = (IH - 1 - r) * GW + c |
| mask[0, 0, r, c] = 1.0 |
|
|
| inits = [ |
| numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'), |
| numpy_helper.from_array(self_idx, 'self_idx'), |
| numpy_helper.from_array(mirror_idx, 'mirror_idx'), |
| numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'), |
| numpy_helper.from_array(mask, 'mask'), |
| ] |
| nodes = [ |
| helper.make_node('Reshape', ['input', 'fs'], ['flat']), |
| helper.make_node('Gather', ['flat', 'self_idx'], ['g_self'], axis=2), |
| helper.make_node('Gather', ['flat', 'mirror_idx'], ['g_mirror'], axis=2), |
| helper.make_node('Add', ['g_self', 'g_mirror'], ['combined']), |
| helper.make_node('Reshape', ['combined', 'os'], ['combined_2d']), |
| helper.make_node('ArgMax', ['combined_2d'], ['am'], axis=1, keepdims=1), |
| ] |
| add_onehot_block(nodes, inits, 'am', 'oh_out') |
| nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output'])) |
| return mk(nodes, inits) |
|
|
| |
| if IW < 2: |
| return None |
| ok = True |
| for inp, out in exs: |
| exp = inp.copy() |
| for r in range(IH): |
| for c in range(IW // 2): |
| v = max(int(inp[r, c]), int(inp[r, IW - 1 - c])) |
| exp[r, c] = v |
| exp[r, IW - 1 - c] = v |
| if not np.array_equal(out, exp): |
| ok = False |
| break |
|
|
| if ok: |
| mirror_idx = np.zeros((GH * GW,), dtype=np.int64) |
| mask = np.zeros((1, 1, GH, GW), dtype=np.float32) |
| self_idx = np.zeros((GH * GW,), dtype=np.int64) |
| for r in range(OH): |
| for c in range(OW): |
| self_idx[r * GW + c] = r * GW + c |
| mirror_idx[r * GW + c] = r * GW + (IW - 1 - c) |
| mask[0, 0, r, c] = 1.0 |
|
|
| inits = [ |
| numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'), |
| numpy_helper.from_array(self_idx, 'self_idx'), |
| numpy_helper.from_array(mirror_idx, 'mirror_idx'), |
| numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'), |
| numpy_helper.from_array(mask, 'mask'), |
| ] |
| nodes = [ |
| helper.make_node('Reshape', ['input', 'fs'], ['flat']), |
| helper.make_node('Gather', ['flat', 'self_idx'], ['g_self'], axis=2), |
| helper.make_node('Gather', ['flat', 'mirror_idx'], ['g_mirror'], axis=2), |
| helper.make_node('Add', ['g_self', 'g_mirror'], ['combined']), |
| helper.make_node('Reshape', ['combined', 'os'], ['combined_2d']), |
| helper.make_node('ArgMax', ['combined_2d'], ['am'], axis=1, keepdims=1), |
| ] |
| add_onehot_block(nodes, inits, 'am', 'oh_out') |
| nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output'])) |
| return mk(nodes, inits) |
|
|
| return None |
|
|
|
|
| def s_extract_inner(td): |
| """Remove N-pixel border frame → smaller output.""" |
| exs = get_exs(td) |
| sp = fixed_shapes(td) |
| if sp is None: |
| return None |
| (IH, IW), (OH, OW) = sp |
|
|
| for b in range(1, min(IH, IW) // 2): |
| if OH != IH - 2 * b or OW != IW - 2 * b: |
| continue |
| if all(np.array_equal(inp[b:IH-b, b:IW-b], out) for inp, out in exs): |
| idx = np.zeros((OH, OW, 2), dtype=np.int64) |
| for r in range(OH): |
| for c in range(OW): |
| idx[r, c] = [r + b, c + b] |
| return _build_gather_model(OH, OW, idx) |
| return None |
|
|
|
|
| def s_add_border(td): |
| """Add constant-color border frame → larger output.""" |
| exs = get_exs(td) |
| sp = fixed_shapes(td) |
| if sp is None: |
| return None |
| (IH, IW), (OH, OW) = sp |
|
|
| for b in range(1, 5): |
| if OH != IH + 2 * b or OW != IW + 2 * b: |
| continue |
| if OH > 30 or OW > 30: |
| continue |
| for bc in range(10): |
| ok = True |
| for inp, out in exs: |
| exp = np.full((OH, OW), bc, dtype=np.int64) |
| exp[b:b+IH, b:b+IW] = inp |
| if not np.array_equal(out, exp): |
| ok = False |
| break |
| if ok: |
| idx = np.zeros((OH, OW, 2), dtype=np.int64) |
| cst = np.full((OH, OW), -1, dtype=np.int64) |
| for r in range(OH): |
| for c in range(OW): |
| if b <= r < b + IH and b <= c < b + IW: |
| idx[r, c] = [r - b, c - b] |
| else: |
| idx[r, c] = [-1, -1] |
| cst[r, c] = bc |
| return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst) |
| return None |
|
|
|
|
| def s_sparse_fill(td): |
| """Each input pixel becomes an NxN block in output.""" |
| exs = get_exs(td) |
| sp = fixed_shapes(td) |
| if sp is None: |
| return None |
| (IH, IW), (OH, OW) = sp |
|
|
| for bh in range(2, 10): |
| for bw in range(2, 10): |
| if OH != IH * bh or OW != IW * bw: |
| continue |
| if OH > 30 or OW > 30: |
| continue |
| ok = True |
| for inp, out in exs: |
| exp = np.zeros((OH, OW), dtype=np.int64) |
| for r in range(IH): |
| for c in range(IW): |
| exp[r*bh:(r+1)*bh, c*bw:(c+1)*bw] = inp[r, c] |
| if not np.array_equal(out, exp): |
| ok = False |
| break |
| if ok: |
| idx = np.zeros((OH, OW, 2), dtype=np.int64) |
| for r in range(OH): |
| for c in range(OW): |
| idx[r, c] = [r // bh, c // bw] |
| return _build_gather_model(OH, OW, idx) |
| return None |
|
|
|
|
| def s_channel_filter(td): |
| """Keep only certain colors, rest → background (0).""" |
| from onnx import helper, numpy_helper |
|
|
| exs = get_exs(td) |
| sp = fixed_shapes(td) |
| if sp is None: |
| return None |
| (IH, IW), (OH, OW) = sp |
| if (IH, IW) != (OH, OW): |
| return None |
|
|
| in_colors = set() |
| out_colors = set() |
| for inp, out in exs: |
| in_colors.update(inp.flatten()) |
| out_colors.update(out.flatten()) |
|
|
| if not (out_colors < in_colors): |
| return None |
|
|
| keep = out_colors |
| for inp, out in exs: |
| exp = np.where(np.isin(inp, list(keep)), inp, 0) |
| if not np.array_equal(out, exp): |
| return None |
|
|
| ch_mask = np.zeros((1, 10, 1, 1), dtype=np.float32) |
| for c in keep: |
| if 0 <= c < 10: |
| ch_mask[0, c, 0, 0] = 1.0 |
|
|
| inits = [numpy_helper.from_array(ch_mask, 'ch_mask')] |
| nodes = [helper.make_node('Mul', ['input', 'ch_mask'], ['output'])] |
| return mk(nodes, inits) |
|
|
