fix: apply SynthSeg's probability per sample, not per batch

[Hendrik-code]

The defect

if self.probability < 1.0 and float(torch.rand((), device=data.device)) >= self.probability:
    return data, target

One draw for the whole batch. probability=0.5 therefore meant "half of all batches are entirely synthetic", not "half of all samples are" — the network saw all-synthetic and all-real batches rather than a mix of both in each. That is a different training signal, and it is out of step with the per-sample convention every other transform in the package follows.

Measured

At B=8, probability=0.5, synthesised rows per batch over five seeds:

before: 8 or 0 only
after:  [5, 4, 4, 6, 4]

Scope

SynthSegTransformsGPU is the standalone driver and carries no registry entry, so no shipped config reaches it — transform_params_paper-synthseg.json uses the registered RandomSynthSegGPU, whose per-sample p comes from kornia and is untouched here. Nothing in the package calls this class outside its own smoke test.

The fix

Draw a per-sample mask, run the generator on the selected rows only, and splice the result back. The unselected samples come out bit-identical rather than merely unchanged in expectation, and the all-selected case still returns the generator's output directly.

Tests

unit_tests/test_synthseg_probability.py: a middling probability must split the batch; 1.0 synthesises everything and 0.0 nothing; unselected rows keep both their image and their labels bit-identical; shape and finiteness hold.

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

🤖 Generated with Claude Code