From 984f037dbf896635c852b9d16605f9bcb7609573 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 15 Oct 2026 09:16:05 +0000 Subject: [PATCH 1/2] feat: give every GPU transfer augmentation a configurable mix range Six of the seventeen transfer augmentations could already blend their output back towards the input, under the name `mix_prob`, from three open-coded copies of the same four lines -- and with the weight drawn from a hardcoded U(0, 1), so the only expressible blends were "none" and "anywhere between the identity and full strength". The other eleven had no blend at all. `mix_alpha` is that weight made configurable: a scalar, or a [lo, hi] range drawn from per sample. It is the weight on the *original*, so 0 is the transform's output and 1 is the input untouched -- the direction the copies used, which is what lets the default of (0, 1) reproduce their draw exactly. `unit_tests/test_mix_prob.py` passes unchanged and is the check that it does. It makes sense for all seventeen: every registered TA is intensity-only on a fixed voxel grid, so none of them resamples, deforms, crops or touches the segmentation, and a per-voxel convex combination is well defined throughout. The geometric transforms are tagged GEO, not TA. `RandomGaussianBlurGPU` comes along because it shares `_RandomConvBaseGPU` and already exposed `mix_prob`. The blend sits before each transform's `retain_stats` restore, where the copies already had it. That is what keeps the promise `retain_stats` makes: blending after the restore would leave the output with the input's mean but a smaller standard deviation, and `mix_prob` is above zero in 29 of the 33 shipped configs, mostly on transforms that ask for `retain_stats` too. One behaviour change, measured rather than assumed. The two per-sample copies spent a draw per sample to decide something `mix_prob` had already settled: if torch.rand(1).item() < self.mix_prob: # drawn even when mix_prob is 0.0 so `RandomInverseGPU` and `RandomHistogramEqualizationGPU` shifted the RNG stream for every transform that ran after them, mixing on or off. Both helpers return before drawing anything now, as the vectorised copy always did. 15 of the 28 runnable shipped configs therefore produce a different seeded sequence, same distribution, and every transform in isolation is bit-identical except `RandomInverseGPU` with `mix_prob > 0`, whose blend also had to move ahead of its hoisted restore. What that buys is the property this feature needs: a `mix_alpha` sweep can be compared against its own `mix_prob = 0` baseline at a fixed seed, which it could not be before. Also here, both found while wiring the above: * `RandomSynthSegGPU.apply_transform` wrote `input[:, c] = ...` without cloning first, mutating the caller's `batch["data"]`. Every sibling already clones; it is uncovered because test_transforms_gpu.py does not list that module. * `AnatomicalLabelOverlay.blend_strength` now shares the range validator, so a value outside [0, 1] is rejected rather than documented. Co-Authored-By: Claude Opus 5 --- smauglab/configs/all_augmentations.json | 83 +++++ smauglab/transforms/gpu/contrast.py | 101 ++++--- smauglab/transforms/gpu/domain_transfer.py | 15 + smauglab/transforms/gpu/fromSeg.py | 26 ++ smauglab/transforms/gpu/mixing.py | 195 ++++++++++++ smauglab/transforms/gpu/palette/overlay.py | 12 +- smauglab/transforms/gpu/palette/transform.py | 16 +- smauglab/transforms/synthseg/transforms.py | 25 ++ unit_tests/test_mix_alpha.py | 301 +++++++++++++++++++ unit_tests/test_mix_prob.py | 6 + 10 files changed, 733 insertions(+), 47 deletions(-) create mode 100644 smauglab/transforms/gpu/mixing.py create mode 100644 unit_tests/test_mix_alpha.py diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 8be9b9b..7301826 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -64,6 +64,11 @@ "keepdim": true, "max_res_aniso": 8.0, "max_res_iso": 4.0, + "mix_alpha": [ + 0.0, + 1.0 + ], + "mix_prob": 0.0, "n_channels": 1, "n_neutral_labels": null, "nonlin_scale": 0.04, @@ -109,6 +114,11 @@ "label_classes": null, "label_remap_prob": 0.5, "min_label_voxels": 4, + "mix_alpha": [ + 0.0, + 1.0 + ], + "mix_prob": 0.0, "n_kmeans_subsample": 10000, "p": 1.0, "p_batch": 1.0, @@ -151,6 +161,11 @@ "dark_threshold": 0.01, "initial_partitioner": null, "keepdim": false, + "mix_alpha": [ + 0.0, + 1.0 + ], + "mix_prob": 0.0, "overlay": null, "p": 1.0, "p_batch": 1.0, @@ -167,6 +182,11 @@ "blend_targets": 1, "include_self": true, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], + "mix_prob": 0.0, "p": 0.2, "p_batch": 1.0, "p_class_mix": 0.0, @@ -186,6 +206,10 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -200,6 +224,10 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -218,6 +246,11 @@ ], "in_seg": 0.2, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], + "mix_prob": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": true, @@ -234,6 +267,10 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -249,6 +286,10 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -263,6 +304,10 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -285,6 +330,10 @@ 5, 7 ], + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -327,6 +376,10 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, @@ -409,7 +462,12 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, + "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, @@ -422,7 +480,12 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, + "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, @@ -435,7 +498,12 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, + "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, @@ -448,7 +516,12 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, + "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, @@ -461,7 +534,12 @@ ], "in_seg": 0.0, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, + "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, @@ -514,7 +592,12 @@ "in_seg": 0.0, "invert": false, "keepdim": true, + "mix_alpha": [ + 0.0, + 1.0 + ], "mix_in_out": false, + "mix_prob": 0.0, "order": 3, "out_seg": 0.0, "p": 1.0, diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index ef6c3c8..e2f9c56 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -10,6 +10,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from +from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, blend_one, configure_mix, mix_with_original from smauglab.transforms.kernels import depthwise_conv3d, gaussian_kernel3d, laplace_kernel, scharr_kernels, stacked_scharr_kernels from smauglab.transforms.rng import shared_choice @@ -59,21 +60,6 @@ def _restore_stats(x: torch.Tensor, stats: tuple[torch.Tensor, torch.Tensor]) -> return (x - new_mean) / (new_std + eps) * orig_stds.view(shape) + orig_means.view(shape) -def _mix_with_original(orig: torch.Tensor, x: torch.Tensor, mix_prob: float) -> torch.Tensor: - """Blend each sample back towards the original with probability `mix_prob`. - - Per sample: an independent draw and an independent alpha for every element of - the batch, which is what "probability of blending the result back" means and - what the sibling transforms that write this inline already do. - """ - if mix_prob <= 0.0: - return x - shape = [x.shape[0]] + [1] * (x.dim() - 1) - mix = torch.rand(x.shape[0], device=x.device).view(shape) < mix_prob - alpha = torch.rand(x.shape[0], device=x.device).view(shape) - return torch.where(mix, alpha * orig + (1 - alpha) * x, x) - - def _check_channel(transform: torch.nn.Module, channel: int, channels: int) -> None: """Reject an `apply_to_channel` entry the input does not have. @@ -258,7 +244,9 @@ class _RandomConvBaseGPU(ImageOnlyTransform): sigma (float): Gaussian width. GaussianBlur and UnsharpMask only. unsharp_amount (float): Strength of the unsharp mask. UnsharpMask only. kernel_sizes (list of int): Multi-scale kernel sizes to draw from. RandConv only. - mix_prob (float): Probability of blending the result back with the original. + mix_prob (float): Per-sample probability of blending the result back with the original. + mix_alpha (float | tuple[float, float]): Weight on the original when blending, fixed or + sampled from a [lo, hi] range. 0 is the convolved output, 1 is the input untouched. retain_stats (bool): If True, restore the original mean and std afterwards. Returns: @@ -283,6 +271,7 @@ def __init__( unsharp_amount: float = 1.0, kernel_sizes: Sequence[int] = (1, 3, 5, 7), mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -301,7 +290,7 @@ def __init__( self.mix_in_out = mix_in_out self.unsharp_amount = unsharp_amount self.kernel_sizes = kernel_sizes - self.mix_prob = mix_prob + configure_mix(self, mix_prob, mix_alpha) def get_kernel(self, device: torch.device) -> Union[Tensor, list[Tensor]]: # Scharr is the odd one out: it returns the three directional kernels as a @@ -411,11 +400,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ else: x = apply_convolution_per_sample(channel_data, self._per_sample_kernels(channel_data.shape[0], input.device)) - # Mix with the original per sample, not once per batch: `mix_prob` is - # documented as the probability of blending the result back with the original, - # a per-sample statement, and RandomInverseGPU / - # RandomHistogramEqualizationGPU run these same three lines inside their loop. - x = _mix_with_original(orig, x, self.mix_prob) + # Before `retain_stats`, not after: the restore is what promises the output + # carries the input's mean and std, and a convex combination of two images + # that already have them does not. See gpu/mixing.py. + x = mix_with_original(self, orig, x) if self.retain_stats: x = _restore_stats(x, stats) @@ -449,6 +437,7 @@ def __init__( self, absolute: bool = False, mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, apply_to_channel: Sequence[int] = (0,), same_on_batch: bool = False, retain_stats: bool = False, @@ -463,6 +452,7 @@ def __init__( kernel_type="Laplace", absolute=absolute, mix_prob=mix_prob, + mix_alpha=mix_alpha, apply_to_channel=apply_to_channel, same_on_batch=same_on_batch, retain_stats=retain_stats, @@ -488,6 +478,7 @@ def __init__( absolute: bool = True, retain_stats: bool = True, mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, apply_to_channel: Sequence[int] = (0,), same_on_batch: bool = False, in_seg: float = 0.0, @@ -502,6 +493,7 @@ def __init__( absolute=absolute, retain_stats=retain_stats, mix_prob=mix_prob, + mix_alpha=mix_alpha, apply_to_channel=apply_to_channel, same_on_batch=same_on_batch, in_seg=in_seg, @@ -531,6 +523,7 @@ def __init__( out_seg: float = 0.0, mix_in_out: bool = False, mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -545,6 +538,7 @@ def __init__( out_seg=out_seg, mix_in_out=mix_in_out, mix_prob=mix_prob, + mix_alpha=mix_alpha, p=p, p_batch=p_batch, keepdim=keepdim, @@ -564,6 +558,7 @@ def __init__( sigma: float = 1.0, unsharp_amount: float = 1.5, mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, apply_to_channel: Sequence[int] = (0,), same_on_batch: bool = False, retain_stats: bool = False, @@ -579,6 +574,7 @@ def __init__( sigma=sigma, unsharp_amount=unsharp_amount, mix_prob=mix_prob, + mix_alpha=mix_alpha, apply_to_channel=apply_to_channel, same_on_batch=same_on_batch, retain_stats=retain_stats, @@ -603,6 +599,7 @@ def __init__( self, kernel_sizes: Sequence[int] = (1, 3, 5, 7), mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, apply_to_channel: Sequence[int] = (0,), same_on_batch: bool = False, retain_stats: bool = False, @@ -617,6 +614,7 @@ def __init__( kernel_type="RandConv", kernel_sizes=kernel_sizes, mix_prob=mix_prob, + mix_alpha=mix_alpha, apply_to_channel=apply_to_channel, same_on_batch=same_on_batch, retain_stats=retain_stats, @@ -1111,6 +1109,11 @@ class _RandomFunctionBaseGPU(ImageOnlyTransform): apply_to_channel (list of int): List of channel indices to apply the function to. Default is [0]. retain_stats (bool): If True, retain the original mean and standard deviation of the image after gamma adjustment. Default is False. same_on_batch (bool): Apply the same transformation across the batch. Default is False. + mix_prob (float): Per-sample probability of blending the result back with the original. + mix_alpha (float | tuple[float, float]): Weight on the original when blending, fixed or + sampled from a [lo, hi] range. 0 is the remapped output, 1 is the input untouched. + Note the function is applied to a per-sample min-max normalisation of the channel, so + with `retain_stats=False` the two sides of the blend are on different scales. p (float): Probability of applying the transform. Default is 1.0. keepdim (bool): Whether to keep the number of dimensions. Default is False. @@ -1127,6 +1130,8 @@ def __init__( in_seg: float = 0.0, out_seg: float = 0.0, mix_in_out: bool = False, + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1138,6 +1143,7 @@ def __init__( self.in_seg = in_seg self.out_seg = out_seg self.mix_in_out = mix_in_out + configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -1163,6 +1169,8 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # Apply function x = self.func(x) + x = mix_with_original(self, orig, x) + if self.retain_stats: x = _restore_stats(x, stats) checked = _select_and_check(self, orig, x, seg_mask) @@ -1204,6 +1212,8 @@ def __init__( in_seg: float = 0.0, out_seg: float = 0.0, mix_in_out: bool = False, + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1216,6 +1226,8 @@ def __init__( in_seg=in_seg, out_seg=out_seg, mix_in_out=mix_in_out, + mix_prob=mix_prob, + mix_alpha=mix_alpha, p=p, p_batch=p_batch, keepdim=keepdim, @@ -1306,6 +1318,7 @@ def __init__( out_seg: float = 0.0, mix_in_out: bool = False, mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1315,8 +1328,8 @@ def __init__( self.retain_stats = retain_stats self.in_seg = in_seg self.out_seg = out_seg - self.mix_prob = mix_prob self.mix_in_out = mix_in_out + configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -1327,14 +1340,26 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ for c in self.apply_to_channel: # The per-sample reductions -- max, and the two retain_stats passes -- are the # expensive part and are independent across samples, so they run once for the - # whole channel. The mix draw and the region selection stay in the loop: + # whole channel. Only the region selection stays in the loop: # `_choose_region_mode` picks `in` / `out` / `all` per sample, and hoisting it - # would make that one decision for the batch. It also keeps the RNG stream. + # would make that one decision for the batch. channel = input[:, c] reduce_dims = tuple(range(1, channel.dim())) keep_shape = (-1, *([1] * (channel.dim() - 1))) inverted = channel.amax(dim=reduce_dims, keepdim=True) - channel + # Before the restore below, which is this transform's `retain_stats`. Vectorised + # rather than drawn per sample inside the loop: the restore is already hoisted, + # and the blend has to precede it. So this transform's seeded output differs + # from the previous release's -- both because the blend moved and because the + # loop it came out of drew a gate per sample even at `mix_prob=0`. See + # gpu/mixing.py for what that cost and why it is gone. + # + # Worth knowing when choosing a range: inverting is `max - x`, so a blend at + # alpha 0.5 is `max / 2` everywhere -- a constant image. The useful ranges sit + # away from the middle. + inverted = mix_with_original(self, channel, inverted) + if self.retain_stats: eps = 1e-8 orig_means = channel.mean(dim=reduce_dims).view(keep_shape) @@ -1347,11 +1372,6 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ orig = channel[i] x = inverted[i] - # Mix with original based on mix_prob - if torch.rand(1).item() < self.mix_prob: - alpha = torch.rand(1, device=input.device) - x = alpha * orig + (1 - alpha) * x - checked = _select_and_check(self, orig, x, None if seg_mask is None else seg_mask[i]) if checked is None: continue @@ -1395,6 +1415,7 @@ def __init__( out_seg: float = 0.0, mix_in_out: bool = False, mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1405,7 +1426,7 @@ def __init__( self.in_seg = in_seg self.out_seg = out_seg self.mix_in_out = mix_in_out - self.mix_prob = mix_prob + configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -1454,12 +1475,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ channel_data = cdf.gather(1, indices).reshape(channel_data.shape).to(channel_data.dtype) - # Mix with original based on mix_prob, per sample. The draws stay in a loop - # so the RNG stream is the one the scalar version produced. + # Per sample, and the draws stay in a loop so the RNG stream is the one the + # scalar version produced -- `mix_prob` is above zero for this transform in most + # of the shipped configs. Before the restore below, as everywhere else. for b in range(batch_size): - if torch.rand(1).item() < self.mix_prob: - alpha = torch.rand(1, device=input.device) - channel_data[b] = alpha * orig[b] + (1 - alpha) * channel_data[b] + channel_data[b] = blend_one(self, orig[b], channel_data[b]) if self.retain_stats: channel_data = _restore_stats(channel_data, stats) @@ -1535,6 +1555,11 @@ class RandomBiasFieldGPU(ImageOnlyTransform): apply_to_channel (list[int]): Channels to which the bias field is applied. invert (bool): If True, uses inverse bias field (1 / field). retain_stats (bool): If True, restores original per-sample mean and std for affected channels. + mix_prob (float): Per-sample probability of blending the result back with the original. + mix_alpha (float | tuple[float, float]): Weight on the original when blending, fixed or + sampled from a [lo, hi] range. 0 is the biased output, 1 is the input untouched. + Blending is close to a rescale of `coefficients` here, since `(1-a) + a*field` is + itself a smooth positive field; it is offered for uniformity across the group. same_on_batch (bool): If True, uses the same sampled coefficients for all batch elements. p (float): Application probability. keepdim (bool): Keep input dimensions flag (passed to base). @@ -1550,6 +1575,8 @@ def __init__( in_seg: float = 0.0, out_seg: float = 0.0, mix_in_out: bool = False, + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, same_on_batch: bool = False, p: float = 1.0, p_batch: float = 1.0, @@ -1571,6 +1598,7 @@ def __init__( self.in_seg = in_seg self.out_seg = out_seg self.mix_in_out = mix_in_out + configure_mix(self, mix_prob, mix_alpha) def _num_coeffs(self, dim: int) -> int: # Count coefficients generated by nested loops matching TorchIO logic. @@ -1685,6 +1713,7 @@ def apply_transform( orig_mean = channel.mean(dim=reduce_dims) orig_std = channel.std(dim=reduce_dims) channel = channel * bias_field + channel = mix_with_original(self, orig, channel) if self.retain_stats: eps = 1e-8 new_mean = channel.mean(dim=reduce_dims) diff --git a/smauglab/transforms/gpu/domain_transfer.py b/smauglab/transforms/gpu/domain_transfer.py index 35d0128..5f972aa 100644 --- a/smauglab/transforms/gpu/domain_transfer.py +++ b/smauglab/transforms/gpu/domain_transfer.py @@ -47,6 +47,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from from smauglab.transforms.gpu.fromSeg import seg_region_masks +from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, blend_at, configure_mix, sample_mix from smauglab.transforms.kernels import gaussian_blur3d, random_bias_field3d # The transfer LUT bank is a multi-hundred-MB artefact built offline by @@ -145,6 +146,8 @@ def __init__( p_spatial_mix: float = 0.0, spatial_mix_scale: float = 0.03, spatial_mix_gain: float = 3.0, + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, same_on_batch: bool = False, p: float = 0.2, p_batch: float = 1.0, @@ -217,6 +220,10 @@ def __init__( self.p_spatial_mix = float(p_spatial_mix) self.spatial_mix_scale = float(spatial_mix_scale) self.spatial_mix_gain = float(spatial_mix_gain) + # Whole-transform blend back towards the input, distinct from `blend_targets` and + # `p_spatial_mix` above: those interpolate between LUTs, which stays a monotone + # remap of the source intensities, while this interpolates in image space. + configure_mix(self, mix_prob, mix_alpha) def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> tuple[Tensor, bool]: """Build the per-class LUTs to use for one sample. @@ -320,6 +327,9 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ N = input.shape[0] for ch in self.apply_to_channel: + # Drawn up front: the per-sample loop below draws its own target LUTs and + # fields, so a blend drawn in place would shift every later sample. + mix_draws = sample_mix(self, N, input.device) for b in range(N): x = input[b, ch] # [D, H, W] — may be z-scored zmode = self._is_zscore(x) @@ -355,6 +365,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ field = _random_bias_field3d(acc.shape, self.bias_field_std, self.bias_scale, acc.device, acc.dtype) acc = (acc * field).clamp(0.0, 1.0) + # Before the z-score restore below, as everywhere else in the group. + # `x01` is the original in the LUT's own [0, 1] domain, which is the space + # `acc` lives in -- blending against the z-scored `x` would mix two scales. + acc = blend_at(mix_draws, b, x01, acc) + if zmode: # restore the input's mean/std (retain-stats): the affine rescale preserves the # transferred distribution SHAPE while placing it back in the z-score scale the diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index a53155a..aac4140 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -9,6 +9,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from +from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, blend_at, configure_mix, mix_with_original, sample_mix from smauglab.transforms.kernels import gaussian_blur3d from smauglab.transforms.rng import shared_choice @@ -278,6 +279,8 @@ def __init__( keepdim: bool = True, std_noise_range: Sequence[float] = (0.1, 0.3), dilation_iterations_range: Sequence[int] = (1, 3), + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, ) -> None: super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.in_seg = in_seg @@ -285,6 +288,7 @@ def __init__( self.retain_stats = retain_stats self.std_noise_range = std_noise_range self.dilation_iterations_range = dilation_iterations_range + configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -344,6 +348,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ dilation_draws = torch.randint( self.dilation_iterations_range[0], self.dilation_iterations_range[1] + 1, (N,), device=input.device ).tolist() + # Hoisted for the same reason as the two above, and for one more: the loop + # below keeps drawing after the blend would land, so drawing in place would + # make switching `mix_prob` on change every later sample's redistribution. + mix_draws = sample_mix(self, N, input.device) # Iterate per sample (seg can differ in shape or labels per sample) for b in range(N): @@ -418,6 +426,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ tmin, tmax = to_add.min(), to_add.max() x = torch.where(tmax - tmin > 1e-8, x + 2 * (to_add - tmin) / (tmax - tmin + 1e-6), x) + # Before the restore below, as everywhere else in the group. `img_batch[b]` + # rather than `input[b, c]` is the original: it is the sanitised channel this + # method works from, and what the no-foreground branch above writes back. + x = blend_at(mix_draws, b, img_batch[b], x) + # Restore stats if self.retain_stats: mean = x.mean() @@ -471,6 +484,11 @@ class RandomPaletteGPU(ImageOnlyTransform): min_label_voxels: minimum voxels in a label to attempt the remap. label_classes: if set, restrict label remap to these class indices (None = all foreground classes present in the batch). + mix_prob: per-sample probability of blending the synthesis back with the original. + mix_alpha: weight on the original when blending, fixed or sampled from a [lo, hi] + range. 0 is the pure synthesis, 1 is the input's own contrast. Both sides of + the blend are in the [0, 1] min-max space the synthesis is built in, and the + foreground z-score below runs after it, so the background stays at zero. p: probability of applying the transform. """ @@ -487,6 +505,8 @@ def __init__( label_remap_prob: float = 0.5, min_label_voxels: int = 4, label_classes: list[int] | None = None, + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, same_on_batch: bool = False, @@ -507,6 +527,7 @@ def __init__( self.label_remap_prob = label_remap_prob self.min_label_voxels = min_label_voxels self.label_classes = label_classes + configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() def apply_transform( @@ -646,6 +667,11 @@ def apply_transform( synth_01 = _gaussian_blur_3d(synth_01, sigma2) synth = synth_01.reshape(B, N) + # Before the foreground z-score below, which is this transform's stats restore. + # Blending here rather than on the way out keeps both sides in the same [0, 1] + # space, and leaves the `* flat_m_all` masking to hold the background at zero. + synth = mix_with_original(self, images_01, synth) + b_sum = (synth * flat_m_all).sum(dim=1, keepdim=True) b_cnt = flat_m_all.sum(dim=1, keepdim=True).clamp(min=1) b_mean = b_sum / b_cnt diff --git a/smauglab/transforms/gpu/mixing.py b/smauglab/transforms/gpu/mixing.py new file mode 100644 index 0000000..7119456 --- /dev/null +++ b/smauglab/transforms/gpu/mixing.py @@ -0,0 +1,195 @@ +"""Blending a transform's output back towards the image it was handed. + +Six of the GPU transfer augmentations already did this, under the name `mix_prob`, +from three open-coded copies of the same four lines -- and with the blend weight drawn +from a hardcoded ``U(0, 1)``, so the only expressible blends were "none" and "anything +from the identity to full strength". `mix_alpha` is that range made configurable, and +this module is the one implementation the whole group shares. + +`alpha` is the weight on the **original**: 0.0 is the transform's output untouched, 1.0 +is the image it was given. That is the direction the open-coded copies used, so the +default range of ``(0.0, 1.0)`` reproduces their draw exactly -- ``draw * 1.0 + 0.0`` is +exact in IEEE-754, and `unit_tests/test_mix_prob.py` is the check that it stayed that +way. + +The blend belongs **before** a transform's `retain_stats` restore, which is where those +copies already sat. `retain_stats` promises that the output carries the input's +per-sample mean and std, and blending after the restore would quietly break that +promise: the convex combination of two images with the same moments has the same mean +but a smaller standard deviation. `mix_prob` is above zero in 29 of the 33 shipped +configs, on transforms that mostly ask for `retain_stats` as well, so that is not a +hypothetical. + +Three entry points, because the transforms come in three shapes: + +* :func:`mix_with_original` for a whole channel, ``[N, *spatial]``, drawing one gate and + one alpha per sample in two vectorised calls. +* :func:`blend_one` for the transforms that assemble their result one sample at a time + inside a Python loop, where a vectorised draw has nothing to vectorise over. +* :func:`sample_mix` with :func:`blend_at` for the loops that *also draw as they go*. + Blending inside one of those would interleave two draws per sample into the transform's + own stream, so switching mixing on would change what every later sample did -- and + `mix_alpha = 0` would stop being a no-op. Drawing up front keeps the blend out of the + way, which is why `gpu/fromSeg.py` already hoists its `in_seg` and dilation draws. + +The two draw differently *on purpose*: each reproduces the draw order of the copy it +replaces, so a transform that already had `mix_prob` set above zero keeps its seeded +output. The one thing that did change is what happens at `mix_prob = 0`. The two +per-sample copies spent a draw per sample to decide something `mix_prob` had already +settled -- + + if torch.rand(1).item() < self.mix_prob: # drawn even when mix_prob is 0.0 + +-- so `RandomInverseGPU` and `RandomHistogramEqualizationGPU` shifted the stream for +every transform that ran after them, whether or not mixing was on. Both helpers return +before drawing anything now, as the vectorised copy always did. Measured consequence: 15 +of the 28 runnable shipped configs produce a different seeded sequence than the previous +release, same distribution; every transform in isolation is bit-identical except +`RandomInverseGPU` with `mix_prob > 0`, whose blend also moved. The property that buys +is the one this feature needs -- a `mix_prob` or `mix_alpha` sweep can be compared +against its own `mix_prob = 0` baseline at a fixed seed, which it could not be before. +""" + +from collections.abc import Sequence +from typing import Protocol, Union + +import torch +from torch import Tensor + +#: A fixed alpha, or a ``[lo, hi]`` range drawn from per sample. Mirrors `BlendSpec` +#: in gpu/palette/overlay.py, which is the same idea for a different blend. +MixAlphaSpec = Union[float, Sequence[float]] + +#: What `mix_alpha` defaults to everywhere: the hardcoded draw this replaced. +DEFAULT_MIX_ALPHA = (0.0, 1.0) + + +class Mixing(Protocol): + """What the helpers below need off the transform they are handed. + + A Protocol rather than `ImageOnlyTransform`, for the same reason `_RegionSelecting` + in gpu/contrast.py is one: these are plain attributes on an `nn.Module`, so reading + them through the base class resolves via `Module.__getattr__`, which is typed as + returning `Tensor | Module` -- and a `Tensor` is not a float. Naming them here is + both what makes the arithmetic below type-check and a statement of the requirement. + """ + + mix_prob: float + mix_alpha: tuple[float, float] + + +#: Appended to a `mix_alpha` bound error, because the direction is the one thing about +#: this parameter nobody guesses right. +_MIX_ALPHA_HINT = "It is the weight on the original input: 0 is the transform's output, 1 is the input untouched." + + +def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str = "", hint: str = _MIX_ALPHA_HINT) -> tuple[float, float]: + """A scalar or ``[lo, hi]`` as a validated pair inside ``[0, 1]``. + + Checked here rather than in `smauglab/config.py`, which only ever asks whether a key + is accepted, never what its value is -- every value check in this package happens in + a constructor. A pipeline is built before the first batch, so a bad range still fails + long before training. + + `key` and `owner` are for the message alone. `AnatomicalLabelOverlay.blend_strength` + is a different blend with the same shape of argument, and shares this. + """ + where = f"{owner}: " if owner else "" + if isinstance(spec, (int, float)) and not isinstance(spec, bool): + lo = hi = float(spec) + else: + try: + values = [float(v) for v in spec] + except (TypeError, ValueError) as exc: + raise ValueError(f"{where}{key} must be a number or a [lo, hi] pair, got {spec!r}") from exc + if len(values) != 2: + raise ValueError(f"{where}{key} must be a number or a [lo, hi] pair, got {len(values)} values: {spec!r}") + lo, hi = values + if lo > hi: + raise ValueError(f"{where}{key} range must be non-decreasing, got [{lo}, {hi}]") + if lo < 0.0 or hi > 1.0: + suffix = f" {hint}" if hint else "" + raise ValueError(f"{where}{key} must lie inside [0, 1], got [{lo}, {hi}].{suffix}") + return (lo, hi) + + +def configure_mix(transform: object, mix_prob: float, mix_alpha: MixAlphaSpec) -> None: + """Validate and assign the two parameters. One line per constructor.""" + transform.mix_prob = float(mix_prob) # type: ignore[attr-defined] + transform.mix_alpha = normalise_blend_range(mix_alpha, owner=type(transform).__name__) # type: ignore[attr-defined] + + +def mix_with_original(transform: Mixing, orig: Tensor, x: Tensor) -> Tensor: + """Blend each sample of a channel back towards `orig`, gated by `mix_prob`. + + `orig` and `x` are ``[N, *spatial]`` -- a whole channel, batch axis first. One + independent gate and one independent alpha per sample, which is what "probability of + blending the result back" means element-wise over a batch. + + Returns `x` untouched, **having drawn nothing**, when `mix_prob` is zero. That matters + beyond the saving: a transform at `mix_prob=0` has to leave the RNG stream exactly + where it found it, or every test that compares a mixed run against an unmixed one at + the same seed would be comparing two different streams. + """ + if transform.mix_prob <= 0.0: + return x + lo, hi = transform.mix_alpha + shape = [x.shape[0]] + [1] * (x.dim() - 1) + # No `dtype=`: the draw this replaced passed none either, so it follows torch's + # default rather than the image's -- which keeps a half-precision batch on the same + # stream as a single-precision one. + mix = torch.rand(x.shape[0], device=x.device).view(shape) < transform.mix_prob + # Drawn even when lo == hi, where the multiply is exactly zero and the result exactly + # `lo`. Skipping it would make the RNG stream a function of `mix_alpha`, so two runs + # that differ only in the range would not be comparable -- which is the one thing a + # configurable range exists to allow. + alpha = torch.rand(x.shape[0], device=x.device).view(shape) * (hi - lo) + lo + return torch.where(mix, alpha * orig + (1 - alpha) * x, x) + + +def sample_mix(transform: Mixing, batch: int, device: torch.device) -> tuple[list[bool], Tensor] | None: + """Every sample's gate and alpha up front, or None when mixing is off. + + The gates come back as a Python list and the alphas stay on the device: reading a + device scalar inside a loop stalls the host until the queue drains, so the one + `.tolist()` here replaces one stall per sample. `gpu/fromSeg.py` hoists its other + per-sample draws the same way and for the same reason. + """ + if transform.mix_prob <= 0.0: + return None + lo, hi = transform.mix_alpha + gates = (torch.rand(batch, device=device) < transform.mix_prob).tolist() + alphas = torch.rand(batch, device=device) * (hi - lo) + lo + return gates, alphas + + +def blend_at(drawn: tuple[list[bool], Tensor] | None, index: int, orig: Tensor, x: Tensor) -> Tensor: + """One sample's blend from a :func:`sample_mix` draw. A no-op when mixing is off.""" + if drawn is None: + return x + gates, alphas = drawn + if not gates[index]: + return x + alpha = alphas[index] + return alpha * orig + (1 - alpha) * x + + +def blend_one(transform: Mixing, orig: Tensor, x: Tensor) -> Tensor: + """The same blend for one sample, out of a Python loop. + + For the transforms that build their result per sample -- a histogram CDF, a + per-region redistribution, a per-sample LUT -- where `orig` and `x` are a single + ``[*spatial]`` volume and there is no batch axis to vectorise over. + + The gate comes off the default generator via `torch.rand(1).item()` and the alpha off + `x`'s device, which is the draw order of the inline copies this replaces; the alpha is + drawn only when the gate fires, as they did. A CPU scalar costs no device + synchronisation, so the loop is not paying for the `.item()`. + """ + if transform.mix_prob <= 0.0: + return x + if torch.rand(1).item() >= transform.mix_prob: + return x + lo, hi = transform.mix_alpha + alpha = torch.rand(1, device=x.device) * (hi - lo) + lo + return alpha * orig + (1 - alpha) * x diff --git a/smauglab/transforms/gpu/palette/overlay.py b/smauglab/transforms/gpu/palette/overlay.py index af2619a..131e3d8 100644 --- a/smauglab/transforms/gpu/palette/overlay.py +++ b/smauglab/transforms/gpu/palette/overlay.py @@ -7,6 +7,7 @@ from torch.nn import functional as F from smauglab.transforms.gpu.fromSeg import foreground_classes +from smauglab.transforms.gpu.mixing import normalise_blend_range #: A scalar blend, or a [lo, hi] range sampled per label per sample. BlendSpec = float | Sequence[float] @@ -52,7 +53,7 @@ def __init__( self.label_remap_prob = float(label_remap_prob) self.min_label_voxels = int(min_label_voxels) self.label_classes = None if label_classes is None else list(label_classes) - self.blend_strength = _normalise_blend(blend_strength) + self.blend_strength = normalise_blend_range(blend_strength, key="blend_strength", owner=type(self).__name__, hint="") self.alpha_magnitude_range = tuple(alpha_magnitude_range) def remap_labels( @@ -103,12 +104,3 @@ def remap_labels( synth = synth * (1.0 - write_mask) + new_vals * write_mask return synth - - -def _normalise_blend(spec: BlendSpec) -> tuple[float, float]: - if isinstance(spec, (int, float)): - return (float(spec), float(spec)) - lo, hi = float(spec[0]), float(spec[1]) - if lo > hi: - raise ValueError(f"blend_strength range must be non-decreasing, got [{lo}, {hi}]") - return (lo, hi) diff --git a/smauglab/transforms/gpu/palette/transform.py b/smauglab/transforms/gpu/palette/transform.py index 61e1326..da16cdb 100644 --- a/smauglab/transforms/gpu/palette/transform.py +++ b/smauglab/transforms/gpu/palette/transform.py @@ -9,6 +9,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from from smauglab.transforms.gpu.fromSeg import _gaussian_blur_3d, collapse_onehot_to_index, voxel_coordinates +from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, configure_mix, mix_with_original from smauglab.transforms.gpu.palette.base import ( BlockContext, InitialPartitioner, @@ -69,6 +70,11 @@ class PaletteSynthesisGPU(ImageOnlyTransform): blur_sigmas_post: the same, for the pass after it. Split in two because a blend below 1.0 wants the overlay's own edges softened, which a single shared draw cannot express. + mix_prob: per-sample probability of blending the synthesis back with the original. + mix_alpha: weight on the original when blending, fixed or sampled from a [lo, hi] + range. 0 is the pure synthesis, 1 is the input's own contrast. Distinct from + the overlay's `blend_strength`, which mixes one label's remap back over the + rest of the synthesis; this one mixes the whole result back over the input. p: probability of applying the transform. """ @@ -81,6 +87,8 @@ def __init__( dark_threshold: float = 0.01, blur_sigmas_pre: Sequence[float] = (0.0, 0.0, 0.0, 0.3, 0.5, 0.8), blur_sigmas_post: Sequence[float] = (0.0, 0.0, 0.0, 0.3, 0.5, 0.8), + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, p_batch: float = 1.0, same_on_batch: bool = False, @@ -95,6 +103,7 @@ def __init__( self.dark_threshold = float(dark_threshold) self.blur_sigmas_pre = list(blur_sigmas_pre) self.blur_sigmas_post = list(blur_sigmas_post) + configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() def apply_transform( @@ -159,7 +168,12 @@ def apply_transform( if sigma2 > 0.0: synth = _gaussian_blur_3d(synth.reshape(B, 1, D, H, W), sigma2).reshape(B, N) - # ── Step 3: foreground z-score ──────────────────────────────────────── + # ── Step 3: blend back towards the input, then foreground z-score ───── + # The z-score is this transform's stats restore, so the blend precedes it, as + # everywhere else in the group. Both sides are in the same [0, 1] space, and the + # `* flat_m_all` below still holds the background at zero. + synth = mix_with_original(self, images_01, synth) + b_cnt = flat_m_all.sum(dim=1, keepdim=True).clamp(min=1) b_mean = (synth * flat_m_all).sum(dim=1, keepdim=True) / b_cnt b_std = (((synth - b_mean) * flat_m_all).pow(2).sum(dim=1, keepdim=True) / b_cnt + eps).sqrt() diff --git a/smauglab/transforms/synthseg/transforms.py b/smauglab/transforms/synthseg/transforms.py index aa14f5e..735cc2b 100644 --- a/smauglab/transforms/synthseg/transforms.py +++ b/smauglab/transforms/synthseg/transforms.py @@ -28,6 +28,7 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from +from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, configure_mix, mix_with_original from smauglab.transforms.synthseg.generator import SynthSegGenerator # Keys understood from the JSON config / kwargs, forwarded to SynthSegGenerator. @@ -102,6 +103,17 @@ class RandomSynthSegGPU(ImageOnlyTransform): Args: apply_to_channel: image channels to overwrite with the synthesis. + mix_prob: per-sample probability of blending the synthesis back with the original. + mix_alpha: weight on the original when blending, fixed or sampled from a [lo, hi] + range. 0 is the pure synthesis, 1 is the real image untouched. + + This is the one transform in the group with no stats restore after the blend: + the synthesis is normalised to [0, 1] while a training patch is z-scored, so + unlike its siblings the two sides are genuinely on different scales. A blend + is still a well-defined convex combination on the same voxel grid -- the + synthesis is intensity-only by construction, see below -- but the result is + not in either input's units. Reach for `RandomPaletteGPU` instead if that + matters, or pair this with `ZscoreNormalizationGPU`. p: application probability (Kornia convention). Remaining kwargs are forwarded to :class:`SynthSegGenerator` (e.g. ``n_channels``, ``prior_means``, ``bias_field_std``, ``gamma_std``, @@ -112,6 +124,8 @@ def __init__( self, apply_to_channel: list[int] | None = None, same_on_batch: bool = False, + mix_prob: float = 0.0, + mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 0.5, p_batch: float = 1.0, keepdim: bool = True, @@ -119,6 +133,10 @@ def __init__( ) -> None: super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.apply_to_channel = apply_to_channel if apply_to_channel is not None else [0] + # Named explicitly rather than left to **kwargs: those are filtered against + # `_GENERATOR_KEYS` on the way to SynthSegGenerator, so a mix parameter arriving + # that way would be dropped without a word. + configure_mix(self, mix_prob, mix_alpha) gen_kwargs = _filter_generator_kwargs(kwargs) # Intensity-only: never deform/flip internally (geometry comes from the # surrounding sequential, which also transports the mask). @@ -131,6 +149,12 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ if seg is None: return input + # A clone, not the caller's tensor: this writes channels back with + # `input[:, c] = ...` and kornia passes the caller's own tensor through when every + # sample applies, which destroyed `batch["data"]` for any caller holding it. Every + # sibling transform already does this; this one did not. + input = input.clone() + # Pass the real image so EM label completion (if enabled) can cluster it. synth, _ = self.generator(seg, image=input) # (B, n_channels, D, H, W) synth = synth.to(device=input.device, dtype=input.dtype) @@ -140,6 +164,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ continue src = n if n < synth.shape[1] else synth.shape[1] - 1 x = synth[:, src] + x = mix_with_original(self, input[:, c], x) if torch.isnan(x).any() or torch.isinf(x).any(): print(f"Warning nan: {self.__class__.__name__}", flush=True) continue diff --git a/unit_tests/test_mix_alpha.py b/unit_tests/test_mix_alpha.py new file mode 100644 index 0000000..78ea788 --- /dev/null +++ b/unit_tests/test_mix_alpha.py @@ -0,0 +1,301 @@ +"""`mix_alpha`: the range a transfer augmentation draws its blend weight from. + +`mix_prob` decides *whether* a sample is blended back towards the image the transform +was handed; `mix_alpha` decides *how far*. Before this existed the weight came from a +hardcoded `U(0, 1)`, so the only expressible blends were "none" and "anywhere between +the identity and full strength" -- and eleven of the seventeen GPU transfer +augmentations had no blend at all. + +alpha is the weight on the **original**: 0.0 is the transform's output, 1.0 is the input +untouched. Counter-intuitive in a config, and kept that way on purpose, because it is the +direction the three open-coded copies used and the default range therefore reproduces +their draw exactly. `test_mix_prob.py` is what holds that. + +The cases that iterate over the registry are the ones that matter most: a transfer +augmentation added later cannot quietly arrive without the knob, or with the knob wired +into its constructor but never read. +""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from smauglab import registry +from smauglab.config import validate_section +from smauglab.registry import AugType, Backend +from smauglab.transforms.build import build_gpu_pipeline +from smauglab.transforms.gpu.mixing import normalise_blend_range +from unit_tests.helpers import SmaugLabTestCase, domain_bank_missing + +BATCH = 6 +SHAPE = (BATCH, 1, 24, 24, 24) + +#: Their closing foreground z-score runs *after* the blend and is unconditional, so +#: alpha=1 hands the z-score the input rather than the synthesis -- which is not the +#: input. Everything else in the group either has no restore or restores the input's own +#: statistics, where that is a no-op. +NOT_IDENTITY_AT_ALPHA_ONE = {"RandomPaletteGPU", "PaletteSynthesisGPU"} + + +def transfer_augmentations(): + """Every registered GPU transfer augmentation.""" + return registry.entries(Backend.GPU, group=AugType.TA) + + +def build(entry, **kwargs): + extra = dict(entry.smoke_kwargs) + extra.update(kwargs) + return entry.cls(p=1.0, **extra) + + +class TestEveryTransferAugmentationHonoursTheKnob(SmaugLabTestCase): + """Registry-driven, so a new TA is covered the moment it is registered.""" + + def volume(self): + return torch.rand(*SHAPE, dtype=torch.float32) + + def seg(self): + seg = torch.zeros(*SHAPE, dtype=torch.float32) + seg[:, :, 6:18, 6:18, 6:18] = 1.0 + seg[:, :, 2:6, 2:6, 2:6] = 2.0 + return seg + + def run_with(self, entry, seed=3, **kwargs): + torch.manual_seed(seed) + transform = build(entry, **kwargs) + torch.manual_seed(seed) + return transform.apply_transform(self.volume().clone(), {"seg": self.seg()}, transform.flags) + + def test_the_constructor_accepts_both_parameters(self): + """The guard against a TA that forgot to declare them.""" + for entry in transfer_augmentations(): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + self.assertIn("mix_prob", accepted) + self.assertIn("mix_alpha", accepted) + + def test_the_defaults_leave_existing_configs_alone(self): + """`mix_prob=0` is off, and `mix_alpha=(0, 1)` is the draw this replaced.""" + for entry in transfer_augmentations(): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + self.assertEqual(accepted["mix_prob"].default, 0.0) + self.assertEqual(tuple(accepted["mix_alpha"].default), (0.0, 1.0)) + + def test_a_mid_range_alpha_changes_the_output(self): + """The guard against a TA whose constructor takes the knob but never reads it.""" + for entry in transfer_augmentations(): + with self.subTest(entry=entry.name): + if entry.external_asset and domain_bank_missing(entry): + self.skipTest(domain_bank_missing(entry)) + plain = self.run_with(entry, mix_prob=0.0) + mixed = self.run_with(entry, mix_prob=1.0, mix_alpha=0.5) + self.assertFalse( + bool(torch.equal(plain, mixed)), + f"{entry.name} accepts mix_alpha but its output does not depend on it", + ) + + def test_alpha_one_is_the_identity(self): + """All the weight on the input means the transform did nothing.""" + for entry in transfer_augmentations(): + with self.subTest(entry=entry.name): + if entry.name in NOT_IDENTITY_AT_ALPHA_ONE: + self.skipTest("an unconditional z-score runs after the blend") + if entry.external_asset and domain_bank_missing(entry): + self.skipTest(domain_bank_missing(entry)) + volume = self.volume() + torch.manual_seed(3) + transform = build(entry, mix_prob=1.0, mix_alpha=1.0) + torch.manual_seed(3) + out = transform.apply_transform(volume.clone(), {"seg": self.seg()}, transform.flags) + # `allclose`, not `equal`: a `retain_stats` restore of the input's own + # mean and std is only the identity up to its own 1e-8 epsilon. + self.assertTrue( + bool(torch.allclose(out, volume, atol=1e-5)), + f"{entry.name} at mix_alpha=1.0 is not the identity (max deviation {float((out - volume).abs().max()):.3g})", + ) + + def test_alpha_zero_is_the_unmixed_transform(self): + """All the weight on the output means the blend did nothing. + + Asserted on `RandomScharrGPU` rather than over the registry, because the + comparison is only sound where the blend's two draws are the *only* randomness + in `apply_transform`. `RandomRedistributeSegGPU` keeps drawing per sample after + the blend, so switching `mix_prob` on necessarily shifts its later samples -- + which is why the draws are hoisted out of its loop rather than left in it, but + hoisting cannot make two extra draws free. `test_alpha_one_is_the_identity` + above is the contract that does hold for every one of them. + """ + from smauglab.transforms.gpu.contrast import RandomScharrGPU + + torch.manual_seed(0) + image = torch.rand(*SHAPE) + mixed = RandomScharrGPU(p=1.0, mix_prob=1.0, mix_alpha=0.0, retain_stats=False) + plain = RandomScharrGPU(p=1.0, mix_prob=0.0, retain_stats=False) + self.assertTrue( + bool( + torch.equal( + mixed.apply_transform(image.clone(), {}, mixed.flags), + plain.apply_transform(image.clone(), {}, plain.flags), + ) + ) + ) + + +class TestTheRangeIsRespected(SmaugLabTestCase): + """Recovering the drawn alpha, as `test_mix_prob.py` does. + + `RandomScharrGPU` is the probe because its kernel is fixed and, with no segmentation + passed, `_choose_region_mode` never fires -- so the mix draws are the only randomness + in `apply_transform` and an exact comparison against an unmixed run is sound. + """ + + def recovered_alphas(self, **kwargs): + from smauglab.transforms.gpu.contrast import RandomScharrGPU + + torch.manual_seed(0) + image = torch.rand(*SHAPE) + + mixed = RandomScharrGPU(p=1.0, mix_prob=1.0, retain_stats=False, **kwargs) + out = mixed.apply_transform(image.clone(), {}, mixed.flags) + plain = RandomScharrGPU(p=1.0, mix_prob=0.0, retain_stats=False) + reference = plain.apply_transform(image.clone(), {}, plain.flags) + + alphas = [] + for b in range(BATCH): + denominator = image[b] - reference[b] + usable = denominator.abs() > 1e-3 + alphas.append(float(((out[b] - reference[b])[usable] / denominator[usable]).median())) + return alphas + + def test_a_narrow_range_bounds_every_draw(self): + alphas = self.recovered_alphas(mix_alpha=[0.8, 0.9]) + for alpha in alphas: + self.assertGreaterEqual(alpha, 0.8 - 1e-3, alphas) + self.assertLessEqual(alpha, 0.9 + 1e-3, alphas) + + def test_two_disjoint_ranges_stay_disjoint(self): + low = self.recovered_alphas(mix_alpha=[0.0, 0.1]) + high = self.recovered_alphas(mix_alpha=[0.9, 1.0]) + self.assertLess(max(low), min(high), f"{low} vs {high}") + + def test_a_scalar_gives_every_sample_the_same_alpha(self): + alphas = [round(a, 4) for a in self.recovered_alphas(mix_alpha=0.35)] + self.assertEqual(set(alphas), {0.35}, alphas) + + def test_the_default_range_still_spreads(self): + """The replacement for the hardcoded U(0, 1) is not accidentally degenerate.""" + alphas = [round(a, 4) for a in self.recovered_alphas()] + self.assertGreater(len(set(alphas)), 1, alphas) + + +class TestMixProbZeroSpendsNothing(SmaugLabTestCase): + """Turning mixing off has to leave the RNG stream where it found it. + + Otherwise a `mix_alpha` sweep could not be compared against its own baseline at a + fixed seed: the transforms downstream would all see a different stream. The two + per-sample copies this replaced got that wrong -- they drew a gate per sample even at + `mix_prob=0` -- and 15 shipped configs changed output when it was fixed. + """ + + def scharr_then_a_draw(self, **kwargs): + from smauglab.transforms.gpu.contrast import RandomScharrGPU + + torch.manual_seed(5) + image = torch.rand(*SHAPE) + transform = RandomScharrGPU(p=1.0, **kwargs) + torch.manual_seed(11) + transform.apply_transform(image.clone(), {}, transform.flags) + return torch.rand(4) + + def test_the_range_does_not_shift_the_stream_when_mixing_is_off(self): + self.assertTrue( + bool(torch.equal(self.scharr_then_a_draw(mix_prob=0.0, mix_alpha=[0.3, 0.3]), self.scharr_then_a_draw(mix_prob=0.0))), + "mix_alpha moved the RNG stream with mix_prob at zero", + ) + + def test_the_range_does_not_shift_the_stream_when_mixing_is_on(self): + """Two draws either way, so an A/B over ranges stays comparable.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(mix_prob=1.0, mix_alpha=[0.2, 0.2]), + self.scharr_then_a_draw(mix_prob=1.0, mix_alpha=[0.4, 0.9]), + ) + ), + "a degenerate range skipped the alpha draw", + ) + + +class TestItSurvivesTheConfigPath(SmaugLabTestCase): + """A range reaches the transform from JSON as a `list`, not a tuple. + + The constructors are the config schema -- `registry.accepted_params` reads + `inspect.signature` -- so a parameter is only really settable once it has been + through `validate_section` and `build_gpu_pipeline`. Asserted on a section rather + than a shipped config file so no new file has to be kept in step. + """ + + SECTION: ClassVar[dict] = { + "RandomScharrGPU": {"p": 1.0, "mix_prob": 0.5, "mix_alpha": [0.2, 0.8]}, + "RandomSynthSegGPU": {"p": 0.0, "mix_alpha": 0.6}, + "RandomLog1pGPU": {"p": 1.0, "mix_prob": 1.0, "mix_alpha": [0.0, 0.25]}, + } + + def test_a_section_validates(self): + self.assertEqual(validate_section(self.SECTION, Backend.GPU), []) + + def test_the_built_transforms_carry_the_parsed_range(self): + built = build_gpu_pipeline(self.SECTION) + by_name = {type(t).__name__: t for t in built} + self.assertEqual(by_name["RandomScharrGPU"].mix_alpha, (0.2, 0.8)) + self.assertEqual(by_name["RandomSynthSegGPU"].mix_alpha, (0.6, 0.6)) + self.assertEqual(by_name["RandomLog1pGPU"].mix_alpha, (0.0, 0.25)) + + def test_a_bad_range_in_a_section_fails_the_build(self): + """Not at validation -- that only checks keys -- but before the first batch.""" + section = {"RandomScharrGPU": {"mix_alpha": [0.9, 0.1]}} + self.assertEqual(validate_section(section, Backend.GPU), []) + with self.assertRaises(ValueError): + build_gpu_pipeline(section) + + +class TestTheRangeIsValidated(SmaugLabTestCase): + """At construction, which is where every value check in this package happens.""" + + def test_a_descending_range_is_rejected(self): + with self.assertRaises(ValueError) as caught: + normalise_blend_range([0.6, 0.4]) + self.assertIn("non-decreasing", str(caught.exception)) + + def test_a_range_outside_the_unit_interval_is_rejected(self): + for bad in ([-0.1, 0.5], [0.5, 1.5], 1.5, -0.1): + with self.subTest(bad=bad), self.assertRaises(ValueError) as caught: + normalise_blend_range(bad) + self.assertIn("[0, 1]", str(caught.exception)) + + def test_a_malformed_range_is_rejected(self): + for bad in ([0.1, 0.2, 0.3], [0.1], "half", [None, 1.0]): + with self.subTest(bad=bad), self.assertRaises(ValueError): + normalise_blend_range(bad) + + def test_the_message_names_the_parameter_and_the_class(self): + from smauglab.transforms.gpu.contrast import RandomScharrGPU + + with self.assertRaises(ValueError) as caught: + RandomScharrGPU(mix_alpha=[0.6, 0.4]) + self.assertIn("mix_alpha", str(caught.exception)) + self.assertIn("RandomScharrGPU", str(caught.exception)) + + def test_a_scalar_and_a_range_both_round_trip(self): + self.assertEqual(normalise_blend_range(0.25), (0.25, 0.25)) + self.assertEqual(normalise_blend_range([0.25, 0.75]), (0.25, 0.75)) + self.assertEqual(normalise_blend_range((0, 1)), (0.0, 1.0)) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/unit_tests/test_mix_prob.py b/unit_tests/test_mix_prob.py index cb529c1..1f7d0a7 100644 --- a/unit_tests/test_mix_prob.py +++ b/unit_tests/test_mix_prob.py @@ -15,6 +15,12 @@ `transform_params_gpu.json` and `transform_params_one-sequence-to-segment-them-all.json` run `RandomScharrGPU` at `mix_prob: 0.5`, so this is a live difference, not a latent one. + +All three copies now live in `smauglab/transforms/gpu/mixing.py`, which also made the +blend weight configurable as `mix_alpha` -- see `test_mix_alpha.py`. This file keeps +working unchanged, which is the point of it: `mix_alpha`'s default is the `U(0, 1)` draw +these cases were written against, so a failure here means that default stopped +reproducing it. """ from __future__ import annotations From 0d7e5ac29b852c6cce18f02f64e6e9fb3b2a0d51 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 15 Oct 2026 09:22:05 +0000 Subject: [PATCH 2/2] fix: make the mix attributes visible to mypy, and reject a bool range Two type errors CI caught that the local mypy (2.1.0 against the pinned 2.4.0) did not. `configure_mix` assigned `mix_prob` / `mix_alpha` from outside the class, so no type checker ever saw them declared and every read resolved through `nn.Module.__getattr__` -- typed `Tensor | Module`, neither of which is a number. That is the exact problem the `Mixing` Protocol was written to solve, and it did not, because structural matching reads the declared attributes. `_RegionSelecting` gets away with the same shape because its three attributes *are* assigned in the constructors. Declaring both on `ImageOnlyTransform`, at the values that mean "off", fixes every read site at once and costs nothing at runtime: only the transforms that take them as constructor arguments can be configured, since the signature is the config schema. `normalise_blend_range` also narrowed `bool` into its sequence branch, where it is not iterable. A bool was never a blend weight anyone meant -- and left to the numeric tower `True` would have quietly meant alpha 1.0, a transform configured to do nothing -- so it is now rejected by name. No behavioural change: every per-transform digest is unchanged, and the blend-before-retain_stats ordering check still holds bit-exactly. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/base.py | 11 +++++++++++ smauglab/transforms/gpu/mixing.py | 18 +++++++++++++----- unit_tests/test_mix_alpha.py | 4 +++- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/smauglab/transforms/gpu/base.py b/smauglab/transforms/gpu/base.py index 65490ad..62f78d3 100644 --- a/smauglab/transforms/gpu/base.py +++ b/smauglab/transforms/gpu/base.py @@ -32,6 +32,8 @@ from torch.distributions import RelaxedBernoulli from torch.nn import Module +from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA + try: # private to kornia, and the only part of the faster parameter draw that needs it from kornia.augmentation.container.image import ImageSequentialBase, _get_new_batch_shape @@ -84,6 +86,15 @@ class ImageOnlyTransform(RigidAffineAugmentationBase3D): """ + #: Blending the result back towards the input; see `gpu/mixing.py`. Declared here, + #: with the values that mean "off", rather than only assigned by `configure_mix` in + #: the constructors that expose them: an attribute a type checker cannot see on the + #: class resolves through `nn.Module.__getattr__`, which is typed `Tensor | Module`, + #: and neither of those is a number. Only the transforms that take them as + #: constructor arguments can be configured, since the signature is the config schema. + mix_prob: float = 0.0 + mix_alpha: tuple[float, float] = DEFAULT_MIX_ALPHA + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: # These transforms never move anything, so the matrix is always the # identity -- shared and cached rather than rebuilt per module. See diff --git a/smauglab/transforms/gpu/mixing.py b/smauglab/transforms/gpu/mixing.py index 7119456..c5fa185 100644 --- a/smauglab/transforms/gpu/mixing.py +++ b/smauglab/transforms/gpu/mixing.py @@ -95,7 +95,11 @@ def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str is a different blend with the same shape of argument, and shares this. """ where = f"{owner}: " if owner else "" - if isinstance(spec, (int, float)) and not isinstance(spec, bool): + if isinstance(spec, bool): # an int subtype, and never a blend weight anyone meant + # ValueError, not the TypeError TRY004 asks for: every other rejection here is a + # ValueError, and a caller should only have to catch one thing for a bad range. + raise ValueError(f"{where}{key} must be a number or a [lo, hi] pair, got {spec!r}") # noqa: TRY004 + if isinstance(spec, (int, float)): lo = hi = float(spec) else: try: @@ -113,10 +117,14 @@ def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str return (lo, hi) -def configure_mix(transform: object, mix_prob: float, mix_alpha: MixAlphaSpec) -> None: - """Validate and assign the two parameters. One line per constructor.""" - transform.mix_prob = float(mix_prob) # type: ignore[attr-defined] - transform.mix_alpha = normalise_blend_range(mix_alpha, owner=type(transform).__name__) # type: ignore[attr-defined] +def configure_mix(transform: Mixing, mix_prob: float, mix_alpha: MixAlphaSpec) -> None: + """Validate and assign the two parameters. One line per constructor. + + `ImageOnlyTransform` declares both attributes, which is what makes `transform` + satisfy `Mixing` here and at every read site -- see the note on that class. + """ + transform.mix_prob = float(mix_prob) + transform.mix_alpha = normalise_blend_range(mix_alpha, owner=type(transform).__name__) def mix_with_original(transform: Mixing, orig: Tensor, x: Tensor) -> Tensor: diff --git a/unit_tests/test_mix_alpha.py b/unit_tests/test_mix_alpha.py index 78ea788..65cdf34 100644 --- a/unit_tests/test_mix_alpha.py +++ b/unit_tests/test_mix_alpha.py @@ -277,7 +277,9 @@ def test_a_range_outside_the_unit_interval_is_rejected(self): self.assertIn("[0, 1]", str(caught.exception)) def test_a_malformed_range_is_rejected(self): - for bad in ([0.1, 0.2, 0.3], [0.1], "half", [None, 1.0]): + # `True` is an int subtype, so without an explicit check it would quietly mean 1.0 + # -- a transform configured to do nothing at all. + for bad in ([0.1, 0.2, 0.3], [0.1], "half", [None, 1.0], True): with self.subTest(bad=bad), self.assertRaises(ValueError): normalise_blend_range(bad)