Skip to content

Hm/deduplicate and declutter - #56

Draft
Hendrik-code wants to merge 3 commits into
mainfrom
hm/deduplicate-and-declutter
Draft

Hm/deduplicate and declutter#56
Hendrik-code wants to merge 3 commits into
mainfrom
hm/deduplicate-and-declutter

Conversation

@Hendrik-code

Copy link
Copy Markdown
Collaborator

Requires #55

Hendrik-code and others added 3 commits August 19, 2026 06:27
… ladders

Which augmentations exist, what they accept, and where they sit in the pipeline
were encoded four times over: one dispatch ladder in gpu/transforms.py, two in
gpu/transforms_list.py, one in cpu/transforms.py, ~900 lines between them. They
had already drifted -- the list pipelines passed a `crop=` argument no transform
accepts, and ordered SimulateLowRes differently from the sequential one.

Each augmentation class now registers itself. The registry (smauglab/registry.py)
is the single source of truth for the class a config key maps to, the parameters
it accepts (read from the constructor signature, so there is no second schema to
drift), its pipeline order, and its GEO/GE/TA group. smauglab/transforms/build.py
does the dispatch once for all three pipeline modes.

Config handling moves to smauglab/config.py: a sectioned GPU/CPU/pipeline schema,
validation that reports every problem in a file at once rather than one per run,
and "did you mean" suggestions. A flat, section-less config is now rejected --
the two namespaces overlapped enough that `GaussianBlurTransform` meant different
transforms depending on which builder read it.

The three nnU-Net trainers collapse into one. Which sections a config populates
decides whether augmentation runs on the dataloader, on the batch, or both, so
the CPU/GPU/Hybrid split no longer needs a class each. The class keeps the name
nnUNetTrainerDAExtGPU: nnU-Net writes it into every checkpoint and resolves the
class from it at inference, so renaming it would strand trained models.

A `smauglab` CLI answers what exists and whether a config is valid by reading the
registry, so it cannot go out of date, and generates the README coverage matrix
and the all-augmentations template config that CI checks for staleness.

The .gitignore ignores *.json repo-wide to keep per-experiment configs out (see
7ff2088). That silently swallowed three things this change adds and needs
tracked -- the generated template, the test fixtures, and the migrator the config
error messages point at -- so each gets an explicit un-ignore.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Each of these is silent -- nothing crashes, nothing fails a test, the pipeline
just produces something other than what the config asked for. Every fix below has
a regression test that fails against the previous implementation.

* RandomFlipTransformGPU never read the flip flags its own generator sampled. It
  recomputed the same `flip_axis`-derived list for every batch element, so it
  flipped all configured axes, identically, on every call -- three seeded calls
  gave byte-identical output, and FlipGenerator3D (including its "at least one
  axis" guarantee) was dead code.

* The 1D Gaussian was sampled at `arange(k)` rather than a centred range, putting
  its peak at index 0. The 3D kernel's maximum sat at corner [0,0,0], so
  RandomGaussianBlurGPU and RandomUnsharpMaskGPU blurred *and* translated the
  image about a voxel -- relative to a segmentation mask that is not convolved.

* in_seg/out_seg reduced the mask's class axis with `argmax(...) > 0`. For an
  ordinary single-channel mask argmax over a length-1 axis is always 0, so the
  result was all-False: in_seg applied the transform nowhere, out_seg applied it
  everywhere. For a one-hot mask it dropped the first foreground class, since
  this repository encodes channel c as label c+1.

* The single-axis generators drew their "random" axis in make_samplers, which
  kornia calls once and caches -- the same axis was degraded for a whole training
  run. CropGenerator3D additionally drew separate axes for the crop and for its
  position, and neutralised the position to 1.0 (the far edge) using the crop's
  neutral value instead of 0.5 (centred).

* The 2D CPU Scharr x-kernel had [-10, 0, -10] as its middle row, summing to -20
  rather than 0. It was not a gradient operator. Predates the registry work.

* The elementwise function transforms normalised with a batch-wide min/max, so a
  volume's augmentation depended on which other volumes shared its batch.

* RandomHistogramEqualizationGPU wrote through an `input[:, c]` view, so its
  non-finite guard `continue`d over values already in the batch.

* RandomChooseXTransformsGPU mutated the caller's batch in place, and raised
  "params must contain 'scale'" for any transform with a kornia parameter
  generator, because it calls apply_transform directly and skips
  forward_parameters. RandomLowResTransformGPU also read flags["data_keys"]
  unguarded, which only the mask path injects.

* Blur sigmas and kernel sizes were drawn with Python's `random`, which
  torch.manual_seed does not reach and which diverges across DDP ranks. They now
  use smauglab.transforms.rng, built from the unused _shared_rand apparatus that
  was already sitting in gpu/fromSeg.py.

Also: scipy's structuring element rank is taken from the data rather than
hardcoded to 3, and `resample_method` is read with .get() so restoring it cannot
raise UnboundLocalError.

Models trained before this change saw the old behaviour and will not reproduce
against it. Configs are unaffected -- no key, parameter or default changed. The
README gains a table of what moved and why.

`smauglab migrate` was advertised in cli.py, config.py and the README but never
existed as a subcommand; those now point at migration/, matching MIGRATE_HINT.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…blocks

Four separate 3D Gaussian blurs, three bias fields and two copies of the
Laplace/Scharr tables had accumulated, each one "kept local so this module stays
self-contained" -- a comment that is still in the diff. They drifted, and the
previous commit had to fix two of them independently. They are now one
smauglab/transforms/kernels.py, promoted from the synthseg implementation, which
was the general case the other three were subsets of.

cpu/artifact.py and cpu/spatial.py held eleven functions between them --
aug_motion, aug_ghosting, aug_spike, aug_bias_field, aug_blur, aug_noise,
aug_swap, aug_flip, aug_affine, aug_elastic, aug_anisotropy -- that were the same
seventeen lines each, differing only in which tio.Random* they constructed. They
become a table and one apply_tio call. That eleven-fold copy had already drifted
once: aug_anisotropy's single-channel branch passed axis=0 to tio.LabelMap, which
nothing else did and which torchio stores as unrecognised metadata.

gpu/contrast.py repeated the same retain-stats / region-select / non-finite-guard
tail in eleven apply_transform methods; the three trainers shared a
character-identical nnU-Net transform tail. Both are extracted. The per-transform
loop structure is deliberately kept rather than inverted into a callback, so the
RNG draw order is provably unchanged.

Three `if __name__ == "__main__":` blocks (~480 lines) go to one parameterised
scripts/demo_augmentations.py. The two GPU copies differed only in a hardcoded
home directory and `cuda(device=7)` vs `cuda()`, so neither could run for anyone
else, and both shipped in the wheel. They also imported cv2, which is not a
dependency; the script writes PNGs through torchvision, which is. Dropping them
removes the only reason for the F811 lint carve-out.

smauglab/utils/utils.py moves to scripts/_common.py: MONAI training-loop helpers,
argparse tuple parsers and a Dice function that nothing under smauglab/ imported
once the demo blocks were gone. config2parser and sig_fn had no callers at all.
The three same-named `normalize` functions, two of which computed different
things, are now normalize_percentile and normalize_minmax.

smauglab/utils/image.py deliberately stays: five modules in the sibling
segtransferaug repository import smauglab.utils.image.Image, so it is public API
in practice despite having no in-package consumer.

CONTRIBUTING gains the rule this commit exists to enforce -- don't write your own
kernel, blur or random draw.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@Hendrik-code Hendrik-code self-assigned this Aug 19, 2026
@Hendrik-code Hendrik-code added the documentation Improvements or additions to documentation label Aug 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant