fix: region selection and per-sample statistics

[Hendrik-code]

Stacked on #58 (which is stacked on #57) — this PR's diff is only its own commit.

Three ways a transform's output depended on something it should not.

in_seg / out_seg did nothing, and the opposite of nothing

_apply_region_mode reduced the mask's class axis with torch.argmax(mask, dim) > 0. That reads "is anything labelled here" only if class 0 is background — and it is not.

• For an ordinary single-channel [B, 1, D, H, W] mask, which is what nnU-Net hands over, 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 collapse_onehot_to_index in gpu/fromSeg.py encodes channel c as label c + 1 with background implicit.

Now amax > 0, extracted as _foreground, which asks the question that was meant and matches what collapse_onehot_to_index already does with seg_raw.any(dim=1).

The function transforms coupled every volume in a batch

_RandomFunctionBaseGPU normalised with a bare x.min() / x.max(), reducing over the whole [N, ...] slab. A volume's augmentation therefore depended on which other volumes happened to share its batch — the same volume in two different batches came out different. Now per-sample amin/amax over dim=1.., as every other transform in the file already does.

aug_redistribute_seg could not handle a 2-D image

The structuring element was built with a hardcoded rank of 3, so scipy raised structure rank must match input rank. Taken from the data instead.

Also, one hazard rather than a fix

RandomHistogramEqualizationGPU now takes input[:, c].clone() rather than the view it wrote back through. This is not observable today and its tests say so. The loop assigns into channel_data[b], so by the time the non-finite guard continued, everything it meant to withhold was already in the batch, and the closing input[:, c] = channel_data was a no-op. Reaching the guard needs the equalisation to produce a non-finite value from finite input, which nothing I could construct manages — a caller-supplied NaN raises out of torch.histc first. The clone makes the guard mean what it says. Same category as the resample change in #57.

Testing

unit_tests/test_region_and_stats.py, 19 tests. Six fail against the parent commit. Because the new _foreground helper would break the import, I verified by reverting only the behavioural call sites and re-running:

FAILED TestApplyRegionMode::test_mode_in_changes_only_the_masked_voxels
FAILED TestApplyRegionMode::test_mode_out_changes_only_the_unmasked_voxels
FAILED TestApplyRegionMode::test_the_unbatched_3d_path_behaves_the_same
FAILED TestRegionModeReachesTheRealTransforms::test_in_seg_confines_a_scharr_transform_to_the_mask
FAILED TestFunctionTransformIsPerSample::test_a_volume_is_augmented_the_same_alone_and_in_a_batch
SUBFAILED(sample=0) TestFunctionTransformIsPerSample::test_each_sample_is_normalised_onto_its_own_range

test_the_old_argmax_reduction_really_did_collapse_it pins the old expression as a control. The end-to-end test drives a Scharr transform through AugmentationSequentialCustom with in_seg=1.0 — the path that actually routes the mask into params["seg"] — and asserts both that nothing outside the mask changed and that something inside it did, which is the half the argmax bug broke.

Full suite: 57 passed, 167 subtests. ruff and mypy clean.

Compatibility

Models trained before this change saw the old behaviour and will not reproduce against it — configs using in_seg/out_seg most of all, since those knobs were inert. No config key, parameter or default changed.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[NathanMolinier]

Looks good