Add SliceJitterTransform: inter-slice misalignment augmentation (#201)

[HanWax]

Summary

Adds SliceJitterTransform to the batchgeneratorsv2 fork, for issue #201 ("any additional augmentations that may improve performance are welcome"). It simulates the inter-slice in-plane misalignment of CT reconstruction — stage drift, sample settling, and sub-volume stitching leave adjacent slices slightly shifted, so geometrically smooth sheets appear jagged/stair-stepped in cross-section. This was floated earlier in the issue thread (as "inter-slice z-jitter") but not implemented.

Model: displacement of slice k along the scan axis = i.i.d. per-slice in-plane shift (uniform, jitter voxels) + smooth low-frequency wander (drift voxels, coarse control points linearly upsampled). Annotations are drawn on the already-misaligned reconstruction, so this is a geometric BasicTransform — image and labels move together (bilinear / nearest), same as Squeeze/Warp.

Implementation: follows the Squeeze/Warp/Decohesion template — randomness sampled once in get_parameters, fully on-device (no CPU/NumPy round-trip), single grid_sample with reflection padding (per the edge-striping finding in #1025), vendored grid helpers for the identical coordinate convention. Regression targets are resampled with bilinear (not the nearest-neighbour path BasicTransform.apply sends them to by default, which would quantize continuous targets) via a small apply() override.

How it differs from existing transforms: elastic deformation is smooth and isotropic — it cannot produce a field that is rigid within a slice yet discontinuous between slices; WarpTransform is the opposite regime (coherent bend, constant along the normal axis). No existing transform produces slice-wise misalignment.

Self-test

python batchgeneratorsv2/transforms/spatial/slice_jitter.py (same convention as the other scroll transforms):

[ok] p_jitter=0 is identity
[ok] shape preserved, finite, no spurious labels
[ok] slice-to-slice misalignment: mean |step|=1.00 vox, range=6.00 vox
[ok] sheets/col median in=25 out=25 (preserved)
[ok] drift-only: max step=0.038 vox, range=6.00 vox (smooth)
[ok] regression target routed through bilinear path
[ok] (160, 160, 160) image+seg: 119.49 ms/sample on cpu

Speed

119 ms/sample at 160³ image+seg on CPU vs 93 ms for the merged WarpTransform in the same environment — same cost class (both are a single grid_sample; sub-ms on GPU).

Status / TODO (draft until complete)

• Before/after figures on a real CT crop (Scroll 5 / PHerc172, matching the validation done for the other Scroll specific 3d augmentations for model training #201 transforms)
• Ablation study (multi-seed Dice, baseline vs augmented) — happy to run on whatever dataset/config you consider the standard setup; guidance welcome
• Defaults / trainer registration per maintainer preference

Closes nothing on its own; contributes to #201.

🤖 Generated with Claude Code

[vercel]

@HanWax is attempting to deploy a commit to the scroll Team on Vercel.

A member of the Team first needs to authorize it.