fix: keep in_seg's redistribution inside the segmentation

[Hendrik-code]

The defect

On the in_seg branch, to_add is written only inside the labels and is exactly zero everywhere else:

if in_seg_bool:
    to_add[l_mask] += ... * (2 * torch.rand(1) - 1)

It was then normalised with a min-max rescale:

img += 2 * (to_add - to_add_min) / (to_add_max - to_add_min + 1e-6)

A min-max rescale moves zero. A voxel where nothing was redistributed picks up -2 * to_add_min / range, which is non-zero whenever to_add_min < 0 — and it always is, because the per-label scale is 2 * rand - 1 on [-1, 1].

So in_seg restricted where the redistribution was computed but not where it landed: the whole patch, background included, got a uniform DC offset. That is the opposite of what the parameter promises.

Measured

16-cube volume, central 8-cube label, in_seg=1.0, comparing against the same image after the function's own [0, 1] normalisation:

before:  max |delta| outside the segmentation = a constant, non-zero offset
after:   max |delta| outside the segmentation = 0.0
         max |delta| inside  the segmentation = 2.0   (unchanged peak)

The fix

Scale by the largest magnitude instead: 2 * to_add / (to_add.abs().max() + 1e-6). Same peak amplitude of 2, and zero maps to zero, so an untouched voxel stays untouched.

Tests

unit_tests/test_redistribute_in_seg.py: nothing changes outside the labels across five seeds; something still changes inside them (the guard against "fixing" this by disabling the transform); the peak amplitude is unchanged; the in_seg=0 whole-image branch still redistributes; and a constant patch stays finite.

Verified failing before the fix (2 failed) and passing after. Full suite unchanged against baseline; ruff check . clean.

🤖 Generated with Claude Code