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/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/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..c5fa185 --- /dev/null +++ b/smauglab/transforms/gpu/mixing.py @@ -0,0 +1,203 @@ +"""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, 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: + 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: 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: + """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..65cdf34 --- /dev/null +++ b/unit_tests/test_mix_alpha.py @@ -0,0 +1,303 @@ +"""`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): + # `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) + + 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