diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 7301826..9da857e 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -41,6 +41,11 @@ "apply_resolution": true, "apply_to_channel": null, "atlas_res": 1.0, + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "bias_field_std": 0.7, "bias_scale": 0.025, "blur_range": 1.03, @@ -61,6 +66,7 @@ "gamma_std": 0.5, "generation_classes": null, "generation_labels": null, + "in_seg": 0.0, "keepdim": true, "max_res_aniso": 8.0, "max_res_iso": 4.0, @@ -74,6 +80,7 @@ "nonlin_scale": 0.04, "nonlin_std": 4.0, "normalise": true, + "out_seg": 0.0, "output_labels": null, "p": 0.5, "p_batch": 1.0, @@ -84,6 +91,9 @@ "rotation_bounds": 15.0, "same_on_batch": false, "scaling_bounds": 0.2, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "shearing_bounds": 0.012, "svf_integration_steps": 7, "thickness": null, @@ -94,6 +104,11 @@ 0.5, 2.0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "blur_sigmas": [ 0.0, 0.0, @@ -110,6 +125,7 @@ 6 ], "dark_threshold": 0.01, + "in_seg": 0.0, "keepdim": false, "label_classes": null, "label_remap_prob": 0.5, @@ -120,6 +136,7 @@ ], "mix_prob": 0.0, "n_kmeans_subsample": 10000, + "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "s_choices": [ @@ -134,6 +151,9 @@ 10 ], "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "skip_parcellation_prob": 0.1, "skip_sub_parc_prob": 0.4 }, @@ -142,6 +162,11 @@ 0.5, 2.0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "blur_sigmas_post": [ 0.0, 0.0, @@ -159,6 +184,7 @@ 0.8 ], "dark_threshold": 0.01, + "in_seg": 0.0, "initial_partitioner": null, "keepdim": false, "mix_alpha": [ @@ -166,20 +192,30 @@ 1.0 ], "mix_prob": 0.0, + "out_seg": 0.0, "overlay": null, "p": 1.0, "p_batch": 1.0, "refinement_partitioners": null, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomDomainTransferGPU": { "any_source": false, "apply_to_channel": null, "bank_path": null, + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "bias_field_std": 0.0, "bias_scale": 0.03, "blend_concentration": 1.0, "blend_targets": 1, + "in_seg": 0.0, "include_self": true, "keepdim": true, "mix_alpha": [ @@ -187,12 +223,16 @@ 1.0 ], "mix_prob": 0.0, + "out_seg": 0.0, "p": 0.2, "p_batch": 1.0, "p_class_mix": 0.0, "p_spatial_mix": 0.0, "pct": 1.0, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "sigma": 2.0, "source_label": null, "spatial_mix_gain": 3.0, @@ -204,47 +244,66 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomHistogramEqualizationGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomRedistributeSegGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", + "confine_prob": 0.2, "dilation_iterations_range": [ 1, 3 ], - "in_seg": 0.2, "keepdim": true, "mix_alpha": [ 0.0, @@ -255,6 +314,9 @@ "p_batch": 1.0, "retain_stats": true, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "label", + "seg_unit_prob": 1.0, "std_noise_range": [ 0.1, 0.3 @@ -265,56 +327,77 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": true, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomLaplaceGPU": { "absolute": false, "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomUnsharpMaskGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "sigma": 1.0, "unsharp_amount": 1.5 }, @@ -322,6 +405,11 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "kernel_sizes": [ @@ -334,217 +422,303 @@ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomClampGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "max_clamp_amount": 0.0, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomGaussianNoiseGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mean": 0.0, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "std": 1.0 }, "RandomGaussianBlurGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "sigma": 1.0 }, "RandomBrightnessGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "brightness_range": [ 0.5, 1.5 ], "in_seg": 0.0, "keepdim": true, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomGammaGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "gamma_range": [ 0.7, 1.5 ], "in_seg": 0.0, "keepdim": false, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomInvGammaGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "gamma_range": [ 0.7, 1.5 ], "in_seg": 0.0, "keepdim": false, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomContrastGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "contrast_range": [ 0.75, 1.25 ], "in_seg": 0.0, "keepdim": true, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomLog1pGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomSqrtGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomSinGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomExpGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomSigmoidGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "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, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomLowResTransformGPU": { "keepdim": true, @@ -588,6 +762,11 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "coefficients": 0.5, "in_seg": 0.0, "invert": false, @@ -596,24 +775,34 @@ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "order": 3, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "ZscoreNormalizationGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "out_seg": 0.0, "p": 1.0, - "p_batch": 1.0 + "p_batch": 1.0, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 } }, "CPU": { @@ -653,7 +842,7 @@ }, "RedistributeTransform": { "classes": null, - "in_seg": 0.2, + "confine_prob": 0.2, "p": 1.0, "retain_stats": false }, diff --git a/smauglab/configs/transform_params.json b/smauglab/configs/transform_params.json index bbaaeaf..d1b5394 100644 --- a/smauglab/configs/transform_params.json +++ b/smauglab/configs/transform_params.json @@ -30,7 +30,7 @@ "retain_stats": true }, "RedistributeTransform": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": true, "p": 0.5 }, diff --git a/smauglab/configs/transform_params_gpu.json b/smauglab/configs/transform_params_gpu.json index 05e170b..493ea1b 100644 --- a/smauglab/configs/transform_params_gpu.json +++ b/smauglab/configs/transform_params_gpu.json @@ -74,7 +74,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.0, "p": 0.25 }, @@ -82,12 +83,13 @@ "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.6, "p": 0.25 }, "RandomRedistributeSegGPU": { - "in_seg": 0.25, + "confine_prob": 0.25, "retain_stats": true, "p": 0.4 }, @@ -97,7 +99,8 @@ "mix_prob": 0.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.25 }, "RandomUnsharpMaskGPU": { @@ -105,7 +108,8 @@ "unsharp_amount": 1.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.25 }, @@ -118,7 +122,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.2 }, @@ -127,7 +132,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomGaussianNoiseGPU": { @@ -135,14 +141,16 @@ "std": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomBrightnessGPU": { @@ -152,7 +160,8 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomGammaGPU": { @@ -163,7 +172,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -174,7 +184,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -185,42 +196,38 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSqrtGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSinGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomExpGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSigmoidGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomLowResTransformGPU": { @@ -256,7 +263,8 @@ "coefficients": 0.2, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "ZscoreNormalizationGPU": { diff --git a/smauglab/configs/transform_params_hybrid.json b/smauglab/configs/transform_params_hybrid.json index 14175bf..db87125 100644 --- a/smauglab/configs/transform_params_hybrid.json +++ b/smauglab/configs/transform_params_hybrid.json @@ -35,18 +35,20 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomHistogramEqualizationGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomRedistributeSegGPU": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.3 }, @@ -55,7 +57,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomUnsharpMaskGPU": { @@ -63,7 +66,8 @@ "unsharp_amount": 1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomRandConvGPU": { @@ -76,7 +80,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.2, "p": 0.1 }, @@ -85,7 +90,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.4 }, "RandomGaussianNoiseGPU": { @@ -93,14 +99,14 @@ "std": 0.1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.1 }, "RandomBrightnessGPU": { @@ -110,7 +116,6 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomGammaGPU": { @@ -121,7 +126,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -132,7 +138,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -143,42 +150,46 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSqrtGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSinGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomExpGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSigmoidGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomLowResTransformGPU": { @@ -202,7 +213,6 @@ "coefficients": 0.5, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.2 }, "ZscoreNormalizationGPU": { @@ -242,7 +252,7 @@ "p": 0.1 }, "RedistributeTransform": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.5 }, diff --git a/smauglab/configs/transform_params_hybrid_TAGE.json b/smauglab/configs/transform_params_hybrid_TAGE.json index 2e2e052..5d62a9b 100644 --- a/smauglab/configs/transform_params_hybrid_TAGE.json +++ b/smauglab/configs/transform_params_hybrid_TAGE.json @@ -36,18 +36,20 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomHistogramEqualizationGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomRedistributeSegGPU": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.3 }, @@ -56,7 +58,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomUnsharpMaskGPU": { @@ -64,7 +67,8 @@ "unsharp_amount": 1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomRandConvGPU": { @@ -77,7 +81,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.2, "p": 0.1 }, @@ -86,7 +91,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.4 }, "RandomGaussianNoiseGPU": { @@ -94,14 +100,14 @@ "std": 0.1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.1 }, "RandomBrightnessGPU": { @@ -111,7 +117,6 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomGammaGPU": { @@ -122,7 +127,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -133,7 +139,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -144,42 +151,46 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSqrtGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSinGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomExpGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSigmoidGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomLowResTransformGPU": { @@ -203,7 +214,6 @@ "coefficients": 0.5, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.2 }, "ZscoreNormalizationGPU": { @@ -243,7 +253,7 @@ "p": 0.1 }, "RedistributeTransform": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.5 }, diff --git a/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json b/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json index 6d124c9..6ed93c6 100755 --- a/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json +++ b/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json @@ -34,7 +34,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.0, "p": 0.25 }, @@ -42,12 +43,13 @@ "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.6, "p": 0.25 }, "RandomRedistributeSegGPU": { - "in_seg": 0.25, + "confine_prob": 0.25, "retain_stats": true, "p": 0.4 }, @@ -57,7 +59,8 @@ "mix_prob": 0.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.25 }, "RandomUnsharpMaskGPU": { @@ -65,7 +68,8 @@ "unsharp_amount": 1.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.25 }, @@ -78,7 +82,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.2 }, @@ -87,7 +92,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomGaussianNoiseGPU": { @@ -95,14 +101,16 @@ "std": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomBrightnessGPU": { @@ -112,7 +120,8 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomGammaGPU": { @@ -123,7 +132,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -134,7 +144,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -145,42 +156,38 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSqrtGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSinGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomExpGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSigmoidGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomLowResTransformGPU": { @@ -204,7 +211,8 @@ "coefficients": 0.2, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "ZscoreNormalizationGPU": { diff --git a/smauglab/configs/transform_params_paper-synthseg.json b/smauglab/configs/transform_params_paper-synthseg.json index dc1b461..f571427 100644 --- a/smauglab/configs/transform_params_paper-synthseg.json +++ b/smauglab/configs/transform_params_paper-synthseg.json @@ -76,7 +76,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.0, "p": 0.0 }, @@ -84,12 +85,13 @@ "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.6, "p": 0.0 }, "RandomRedistributeSegGPU": { - "in_seg": 0.25, + "confine_prob": 0.25, "retain_stats": true, "p": 0.0 }, @@ -99,7 +101,8 @@ "mix_prob": 0.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomUnsharpMaskGPU": { @@ -107,7 +110,8 @@ "unsharp_amount": 1.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.0 }, @@ -120,7 +124,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.0 }, @@ -129,7 +134,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomGaussianNoiseGPU": { @@ -137,14 +143,16 @@ "std": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomBrightnessGPU": { @@ -154,7 +162,8 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomGammaGPU": { @@ -165,7 +174,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -176,7 +186,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -187,42 +198,38 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSqrtGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSinGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomExpGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomSigmoidGPU": { "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": false, "p": 0.1 }, "RandomLowResTransformGPU": { @@ -246,7 +253,8 @@ "coefficients": 0.2, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "ZscoreNormalizationGPU": { diff --git a/smauglab/registry.py b/smauglab/registry.py index 36805bc..4dd7351 100644 --- a/smauglab/registry.py +++ b/smauglab/registry.py @@ -177,6 +177,16 @@ def __init__(self, source: str, problems: list[str]) -> None: "kernel_type": "the class name now carries the kernel (e.g. RandomScharrGPU)", "func": "the class name now carries the function (e.g. RandomLog1pGPU)", "one_dim": "use RandomAcqTransformGPU (one_dim) or RandomLowResTransformGPU", + "mix_in_out": "use seg_unit='label' with seg_unit_prob", + # Only reachable on RandomRedistributeSegGPU, the one transform that does not + # accept `in_seg`. That is deliberate: it used to accept it with a *different* + # meaning (confine each region's Gaussian to its own voxels), twenty-eight shipped + # configs set it, and if the class also took the standard `in_seg` they would all + # have kept validating while silently changing what they do. Keeping the name + # unaccepted is what turns that into this message. Note difflib would otherwise + # answer "did you mean out_seg?" -- it scores 0.615, above the 0.6 cutoff -- which + # is exactly why `out_seg` is not accepted there either. + "in_seg": "on RandomRedistributeSegGPU it is now confine_prob; elsewhere in_seg is unchanged", } ) diff --git a/smauglab/transforms/cpu/fromSeg.py b/smauglab/transforms/cpu/fromSeg.py index 6f467a2..11e64db 100644 --- a/smauglab/transforms/cpu/fromSeg.py +++ b/smauglab/transforms/cpu/fromSeg.py @@ -19,14 +19,18 @@ class RedistributeTransform(BasicTransform): Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py """ - def __init__(self, classes=None, in_seg=0.2, retain_stats=False): + # `confine_prob` was `in_seg`, renamed with its GPU twin: it is the probability that + # each label's Gaussian is confined to that label's own voxels, not the probability of + # applying the transform inside the segmentation, which is what `in_seg` means + # everywhere else. Both halves of `AugId.REDISTRIBUTE_SEG` use the one name. + def __init__(self, classes=None, confine_prob=0.2, retain_stats=False): super().__init__() self.classes = classes - self.in_seg = in_seg + self.confine_prob = confine_prob self.retain_stats = retain_stats def get_parameters(self, **data_dict) -> dict: - return {"classes": self.classes, "in_seg": self.in_seg, "retain_stats": self.retain_stats} + return {"classes": self.classes, "confine_prob": self.confine_prob, "retain_stats": self.retain_stats} def apply(self, data_dict: dict, **params) -> dict: if data_dict.get("image") is not None and data_dict.get("segmentation") is not None: @@ -36,12 +40,12 @@ def apply(self, data_dict: dict, **params) -> dict: def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tuple[torch.Tensor, torch.Tensor]: for c in range(1): # Works on the first channel only img[c], seg[c] = aug_redistribute_seg( - img[c], seg[c], classes=params["classes"], in_seg=params["in_seg"], retain_stats=params["retain_stats"] + img[c], seg[c], classes=params["classes"], confine_prob=params["confine_prob"], retain_stats=params["retain_stats"] ) return img, seg -def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False): +def aug_redistribute_seg(img, seg, classes=None, confine_prob=0.2, retain_stats=False): """ Augment the image by redistributing the values of the image within the regions defined by the segmentation. @@ -50,7 +54,7 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) """ device = img.device _seg = seg - in_seg_bool = 1 - torch.rand(1, device=device) <= in_seg + confine = 1 - torch.rand(1, device=device) <= confine_prob if classes: _seg = combine_classes(_seg, classes) @@ -112,7 +116,7 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) ) sign = 2 * torch.rand(1, device=device) - 1 - if in_seg_bool: + if confine: to_add[l_mask] += _normal_pdf(img[l_mask], l_mean, redist_std) * sign else: to_add += _normal_pdf(img, l_mean, redist_std) * sign @@ -120,8 +124,8 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) # The min-max form shifted zero: with `2 * (to_add - min) / range`, a voxel where # nothing was redistributed picks up `-2 * min / range`, non-zero whenever min < 0 -- # and it always is, since the per-label scale is `2 * rand - 1` in [-1, 1]. On the - # `in_seg` branch `to_add` is exactly zero outside the labels, so that constant became - # a DC offset over the whole patch: `in_seg` restricted where the redistribution was + # confined branch `to_add` is exactly zero outside the labels, so that constant became + # a DC offset over the whole patch: confining restricted where the redistribution was # *computed* but not where it landed. # # Scaling by the largest magnitude keeps the peak amplitude of 2 and maps zero to zero, diff --git a/smauglab/transforms/gpu/base.py b/smauglab/transforms/gpu/base.py index 62f78d3..d43e55a 100644 --- a/smauglab/transforms/gpu/base.py +++ b/smauglab/transforms/gpu/base.py @@ -33,6 +33,13 @@ from torch.nn import Module from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, +) 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 @@ -95,6 +102,24 @@ class ImageOnlyTransform(RigidAffineAugmentationBase3D): mix_prob: float = 0.0 mix_alpha: tuple[float, float] = DEFAULT_MIX_ALPHA + #: Which part of the segmentation the transform applies to, and how sharply it stops; + #: see `gpu/regions.py`. Declared here for the same reason as the two above, and with + #: the values that mean "off": `in_seg = out_seg = 0` is "the whole patch", and then + #: none of the other five is read. `seg_feather` is the normalised pair + #: `configure_regions` assigns, not the scalar-or-range a config may write. + #: + #: These used to be assigned only in the constructors, with no class-level default, so + #: `ZscoreNormalizationGPU` had to set `mix_in_out` by hand outside its signature to + #: keep the helper's contract -- the registry derives a config's accepted keys from + #: `__init__`, so declaring it there would have invented a setting. + in_seg: float = 0.0 + out_seg: float = 0.0 + seg_unit: str = DEFAULT_SEG_UNIT + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB + bg_split: str = DEFAULT_BG_SPLIT + bg_em_clusters: tuple[int, int] = DEFAULT_BG_EM_CLUSTERS + seg_feather: tuple[float, float] = (DEFAULT_SEG_FEATHER, DEFAULT_SEG_FEATHER) + 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 e2f9c56..23321eb 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -1,7 +1,7 @@ import functools import math from collections.abc import Callable, Sequence -from typing import Any, Protocol, Union +from typing import Any, Union import torch import torchvision.transforms._functional_tensor as F_t @@ -11,33 +11,25 @@ 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 + +# Re-exported, not used here: `_foreground` moved to gpu/regions.py with the rest of the +# region machinery, and the tests that pin its three segmentation layouts import it from +# this module, where it used to live. +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _foreground, # noqa: F401 + _select_and_check, + configure_regions, +) from smauglab.transforms.kernels import depthwise_conv3d, gaussian_kernel3d, laplace_kernel, scharr_kernels, stacked_scharr_kernels from smauglab.transforms.rng import shared_choice -def _choose_region_mode( - p_in: float, - p_out: float, - seg_mask: torch.Tensor | None, # noqa: ARG001 -) -> str: - """Sample where to apply the transform: 'in', 'out', or 'all'. - - - p_in, p_out are probabilities in [0,1]. - - If both probs are 0, or both fire at once, return 'all'. - - seg_mask is accepted but unused here; _apply_region_mode treats a None - mask as 'all' regardless of the mode chosen. - """ - p_in = float(max(0.0, min(1.0, p_in))) - p_out = float(max(0.0, min(1.0, p_out))) - in_bool = torch.rand(()) < p_in - out_bool = torch.rand(()) < p_out - if in_bool and not out_bool: - return "in" - if out_bool and not in_bool: - return "out" - return "all" - - def _channel_stats(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-sample mean and std of a `[N, ...spatial]` channel, each shaped `[N]`.""" reduce_dims = tuple(range(1, x.dim())) @@ -76,159 +68,81 @@ def _check_channel(transform: torch.nn.Module, channel: int, channels: int) -> N ) -class _RegionSelecting(Protocol): - """What `_select_and_check` needs off the transform it is handed. +def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torch.Tensor: + """ + Based on https://github.com/pytorch/vision/blob/e3b5d3a8bf5e8636462fd8bce9897bccc690b2a0/torchvision/transforms/_functional_tensor.py#L746 + """ + if not (isinstance(img, torch.Tensor)): + raise TypeError(f"img should be Tensor. Got {type(img)}") + + if dim == 2: + kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1]) + padding = [kernel.shape[2] // 2, kernel.shape[2] // 2, kernel.shape[3] // 2, kernel.shape[3] // 2] + elif dim == 3: + kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1], kernel.shape[2]) + padding = [ + kernel.shape[2] // 2, + kernel.shape[2] // 2, + kernel.shape[3] // 2, + kernel.shape[3] // 2, + kernel.shape[4] // 2, + kernel.shape[4] // 2, + ] + else: + raise ValueError(f"Only 2D and 3D convolution are supported. Got {dim}D.") - A Protocol rather than `ImageOnlyTransform`, because these are plain attributes on - an nn.Module: reading them through the base class resolves via `Module.__getattr__`, - which is typed as returning `Tensor | Module`. Naming them here is both what makes - that type-check and a statement of the helper's actual requirement. + img, need_cast, need_squeeze, out_dtype = F_t._cast_squeeze_in(img, [kernel.dtype]) - All three are attributes rather than constructor parameters as far as this - Protocol is concerned; the registry derives a transform's config surface from its - __init__ signature, so declaring one here does not make it settable in a config. - """ + # padding = (left, right, top, bottom) + img = F.pad(img, padding, mode="reflect") + if dim == 2: # noqa: SIM108 -- the 2d/3d split reads better spelled out than as a ternary + img = F.conv2d(img, kernel, groups=img.shape[-(1 + dim)]) + else: # dim == 3 + # Via `depthwise_conv3d`, not `F.conv3d`: every row of `kernel` is the same + # filter (it came from an `expand` above), so one row is handed over and the + # helper broadcasts it back across the planes, which is also what lets it make + # `groups > 1` from a single-plane call. The bucketed pipelines hand over one + # sample at a time, precisely the case cuDNN serves 35x slower than depthwise. + leading = img.shape[:-3] + planes = int(torch.tensor(leading).prod()) if leading else 1 + out = depthwise_conv3d(img.reshape(planes, 1, *img.shape[-3:]), kernel[:1, :1]) + img = out.reshape(*leading, *out.shape[-3:]) - in_seg: float - out_seg: float - mix_in_out: bool + img = F_t._cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) + return img -def _select_and_check( - transform: _RegionSelecting, - orig: torch.Tensor, - x: torch.Tensor, - seg_mask: torch.Tensor | None, - note: str = "", -) -> torch.Tensor | None: - """Apply region selection, then reject the result if it went non-finite. +def apply_convolution_per_sample(channel_data: torch.Tensor, kernels: Sequence[Tensor]) -> torch.Tensor: + """Convolve `[N, D, H, W]` with one 3-D kernel per sample, in a single conv. - Returns None when the channel should be left as it was -- the "Final safety" check - that closed all nine of these loops identically. + The alternative -- and what every caller here used to do -- is a Python loop + calling `apply_convolution` on `channel_data[b : b + 1]`. That is a + single-channel `F.conv3d`, which cuDNN serves from its implicit-GEMM path: 4.3 ms + per sample for a 3x3x3 kernel over a 128^3 patch on an A40, against 0.12 ms for + the whole batch here. See `kernels.depthwise_conv3d` for why the shape matters. - Takes the transform rather than its three region attributes: every call site read - exactly `self.in_seg`, `self.out_seg` and `self.mix_in_out`, and spelling them out - made the call longer than the code it replaces. - """ - if seg_mask is not None: - region_mode = _choose_region_mode(transform.in_seg, transform.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=transform.mix_in_out) - # One `isfinite(...).all()`, not `isnan().any() or isinf().any()`: each of those - # reads a reduction back to the host, and the `or` forces the second on the common - # path. Runs for every channel of every region-selecting transform. - if not bool(torch.isfinite(x).all()): - print(f"Warning nan: {type(transform).__name__}{note}", flush=True) - return None - return x - - -def _foreground(mask: torch.Tensor, dim: int) -> torch.Tensor: - """Which voxels the segmentation covers, reducing over the class axis. - - This used to be `torch.argmax(mask, dim) > 0`, which is only "is anything labelled - here" if class 0 is background -- and it is not: - - * For a single-channel `[B, 1, D, H, W]` mask (an ordinary nnU-Net target, and what - the tests build) `argmax` over a length-1 axis is always 0, so the result was - **all False**. `in_seg` then applied the transform nowhere and `out_seg` applied - it everywhere: the two knobs did nothing and the opposite of nothing. - * For a one-hot mask, this repository's convention (`collapse_onehot_to_index` in - gpu/fromSeg.py) is that channel `c` encodes label `c + 1` with background - implicit, so `argmax == 0` is a real foreground class and was being dropped. - - `amax > 0` asks the question that was meant, and matches what - `collapse_onehot_to_index` already does with `seg_raw.any(dim=1)`. - - That still leaves one layout it cannot read off the values alone: a one-hot - that *includes* a background channel, which is what `seg_region_masks` emits - for a single-channel label map and what `unit_tests/test_seg_layout.py` - builds. Every voxel then has some channel set, `amax > 0` is True everywhere, - and `in_seg` applies the transform to the whole patch while `out_seg` applies - it nowhere -- the exact failure this function was written to fix, in the one - layout it was not checked against. - - A background-inclusive one-hot is recognisable: it is multi-channel and leaves - no voxel unset, because the background channel covers whatever the foreground - channels do not. A background-implicit one-hot always has all-zero background - voxels unless the patch is labelled edge to edge, which a real segmentation - patch is not. So full coverage plus more than one channel means channel 0 is - background, and the reduction skips it. - """ - covered = mask.amax(dim=dim) > 0 - if mask.shape[dim] == 1: - return covered - # `torch.where`, not `if bool(...)`: reading that test back to Python stalls the host - # on everything queued behind it. Both branches are a reduction over the class axis, - # so computing the unused one costs a fraction of what the stall does. - without_background = mask.narrow(dim, 1, mask.shape[dim] - 1).amax(dim=dim) > 0 - return torch.where(covered.all(), without_background, covered) - - -def _apply_region_mode( - orig: torch.Tensor, - transformed: torch.Tensor, - seg_mask: torch.Tensor | None, - mode: str, - normalize: bool = False, - mix_in_out: bool = False, -) -> torch.Tensor: - """Blend transformed with orig based on region selection mode. - - - mode 'all': return transformed - - mode 'in': apply transform inside seg, keep orig outside - - mode 'out': apply transform outside seg, keep orig inside - - mix_in_out: if True, randomly apply transform to some of the segmentation, not all. + Kernels of different sizes are zero-padded up to the widest one. That is exact, + not an approximation: a zero tap contributes nothing, and `reflect` padding gives + the same value at a given offset however wide the pad is, so the extra taps read + real (if irrelevant) voxels and multiply them by zero. `RandConv` draws its kernel + size per sample, so without this the batch could not be convolved in one call. """ - if seg_mask is None or mode == "all": - return transformed - - # Rescale transformed based on min max orig - # Needed due to the important change in the image - if orig.dim() == 4: - if normalize: - orig_min = torch.amin(orig, dim=tuple(range(1, orig.dim())), keepdim=True) - orig_max = torch.amax(orig, dim=tuple(range(1, orig.dim())), keepdim=True) - transformed_min = torch.amin(transformed, dim=tuple(range(1, transformed.dim())), keepdim=True) - transformed_max = torch.amax(transformed, dim=tuple(range(1, transformed.dim())), keepdim=True) - transformed = (transformed - transformed_min) / (transformed_max - transformed_min + 1e-8) * (orig_max - orig_min) + orig_min - - # No `.clone()`: nothing below writes into `m`, and on a [B, C, D, H, W] mask - # that copy is the size of the batch. - m = seg_mask.to(transformed.dtype) - if mix_in_out: - # One keep/drop draw per (sample, class), in a single call rather than one - # `randint` per sample: the same draws in the same order, one kernel. - keep = torch.randint(0, 2, seg_mask.shape[:2], device=seg_mask.device, dtype=m.dtype) - m = m * keep.view(*seg_mask.shape[:2], *([1] * (m.dim() - 2))) - - m = _foreground(m, dim=1) - m = m.to(transformed.dtype) - if mode == "out": - m = 1.0 - m - - elif orig.dim() == 3: - if normalize: - orig_min = torch.amin(orig) - orig_max = torch.amax(orig) - transformed_min = torch.amin(transformed) - transformed_max = torch.amax(transformed) - transformed = (transformed - transformed_min) / (transformed_max - transformed_min + 1e-8) * (orig_max - orig_min) + orig_min - - m = seg_mask.to(transformed.dtype) - if mix_in_out: - # Create a tensor with random one and zero - o = torch.randint(0, 2, (seg_mask.shape[0],), device=seg_mask.device, dtype=m.dtype) - m = m * o.view(-1, 1, 1, 1) # Broadcasting o to match the dimensions of m - m = _foreground(m, dim=0) - m = m.to(transformed.dtype) - if mode == "out": - m = 1.0 - m - - else: - raise ValueError(f"Only 4D and 3D images are supported. Got {orig.dim()}D.") + sizes = [int(k.shape[-1]) for k in kernels] + widest = max(sizes) + stacked = [] + for kernel, size in zip(kernels, sizes): + if size == widest: + stacked.append(kernel) + else: + lo = (widest - size) // 2 + padded = torch.zeros((widest, widest, widest), device=kernel.device, dtype=kernel.dtype) + padded[lo : lo + size, lo : lo + size, lo : lo + size] = kernel + stacked.append(padded) + weight = torch.stack(stacked).unsqueeze(1) # [N, 1, k, k, k] - return m * transformed + (1.0 - m) * orig + pad = widest // 2 + padded_img = F.pad(channel_data.unsqueeze(1), [pad] * 6, mode="reflect") + return depthwise_conv3d(padded_img, weight).squeeze(1) ## Convolution transform @@ -262,7 +176,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, # Kernel-specific, and explicit rather than read from **kwargs, where # `inspect.signature` could not see them and a config typo silently took the # default. Defaults are the historical kwargs.get() ones, so behaviour is unchanged. @@ -285,9 +203,7 @@ def __init__( self.absolute = absolute self.sigma = sigma self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) self.unsharp_amount = unsharp_amount self.kernel_sizes = kernel_sizes configure_mix(self, mix_prob, mix_alpha) @@ -443,7 +359,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -458,7 +378,11 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -483,7 +407,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -498,7 +426,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -521,7 +453,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -536,7 +472,11 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, mix_prob=mix_prob, mix_alpha=mix_alpha, p=p, @@ -564,7 +504,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -580,7 +524,11 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -605,7 +553,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -620,90 +572,17 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, ) -def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torch.Tensor: - """ - Based on https://github.com/pytorch/vision/blob/e3b5d3a8bf5e8636462fd8bce9897bccc690b2a0/torchvision/transforms/_functional_tensor.py#L746 - """ - if not (isinstance(img, torch.Tensor)): - raise TypeError(f"img should be Tensor. Got {type(img)}") - - if dim == 2: - kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1]) - padding = [kernel.shape[2] // 2, kernel.shape[2] // 2, kernel.shape[3] // 2, kernel.shape[3] // 2] - elif dim == 3: - kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1], kernel.shape[2]) - padding = [ - kernel.shape[2] // 2, - kernel.shape[2] // 2, - kernel.shape[3] // 2, - kernel.shape[3] // 2, - kernel.shape[4] // 2, - kernel.shape[4] // 2, - ] - else: - raise ValueError(f"Only 2D and 3D convolution are supported. Got {dim}D.") - - img, need_cast, need_squeeze, out_dtype = F_t._cast_squeeze_in(img, [kernel.dtype]) - - # padding = (left, right, top, bottom) - img = F.pad(img, padding, mode="reflect") - if dim == 2: # noqa: SIM108 -- the 2d/3d split reads better spelled out than as a ternary - img = F.conv2d(img, kernel, groups=img.shape[-(1 + dim)]) - else: # dim == 3 - # Via `depthwise_conv3d`, not `F.conv3d`: every row of `kernel` is the same - # filter (it came from an `expand` above), so one row is handed over and the - # helper broadcasts it back across the planes, which is also what lets it make - # `groups > 1` from a single-plane call. The bucketed pipelines hand over one - # sample at a time, precisely the case cuDNN serves 35x slower than depthwise. - leading = img.shape[:-3] - planes = int(torch.tensor(leading).prod()) if leading else 1 - out = depthwise_conv3d(img.reshape(planes, 1, *img.shape[-3:]), kernel[:1, :1]) - img = out.reshape(*leading, *out.shape[-3:]) - - img = F_t._cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) - return img - - -def apply_convolution_per_sample(channel_data: torch.Tensor, kernels: Sequence[Tensor]) -> torch.Tensor: - """Convolve `[N, D, H, W]` with one 3-D kernel per sample, in a single conv. - - The alternative -- and what every caller here used to do -- is a Python loop - calling `apply_convolution` on `channel_data[b : b + 1]`. That is a - single-channel `F.conv3d`, which cuDNN serves from its implicit-GEMM path: 4.3 ms - per sample for a 3x3x3 kernel over a 128^3 patch on an A40, against 0.12 ms for - the whole batch here. See `kernels.depthwise_conv3d` for why the shape matters. - - Kernels of different sizes are zero-padded up to the widest one. That is exact, - not an approximation: a zero tap contributes nothing, and `reflect` padding gives - the same value at a given offset however wide the pad is, so the extra taps read - real (if irrelevant) voxels and multiply them by zero. `RandConv` draws its kernel - size per sample, so without this the batch could not be convolved in one call. - """ - sizes = [int(k.shape[-1]) for k in kernels] - widest = max(sizes) - stacked = [] - for kernel, size in zip(kernels, sizes): - if size == widest: - stacked.append(kernel) - else: - lo = (widest - size) // 2 - padded = torch.zeros((widest, widest, widest), device=kernel.device, dtype=kernel.dtype) - padded[lo : lo + size, lo : lo + size, lo : lo + size] = kernel - stacked.append(padded) - weight = torch.stack(stacked).unsqueeze(1) # [N, 1, k, k, k] - - pad = widest // 2 - padded_img = F.pad(channel_data.unsqueeze(1), [pad] * 6, mode="reflect") - return depthwise_conv3d(padded_img, weight).squeeze(1) - - ## Noise transform @register( aug_id=AugId.GAUSSIAN_NOISE, @@ -733,7 +612,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -742,9 +625,7 @@ def __init__( self.apply_to_channel = apply_to_channel self.mean = mean self.std = std - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @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: @@ -801,7 +682,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -809,9 +694,7 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.brightness_range = brightness_range self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @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: @@ -870,7 +753,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = False, @@ -880,9 +767,7 @@ def __init__( self.invert_image = invert_image self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @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: @@ -963,7 +848,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = False, @@ -976,7 +865,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -999,7 +892,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = False, @@ -1012,7 +909,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -1049,7 +950,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1058,9 +963,7 @@ def __init__( self.contrast_range = contrast_range self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @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: @@ -1129,7 +1032,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1140,9 +1047,7 @@ def __init__( self.func = func self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency @@ -1211,7 +1116,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1225,7 +1134,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, mix_prob=mix_prob, mix_alpha=mix_alpha, p=p, @@ -1316,7 +1229,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1326,9 +1243,7 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency @@ -1340,9 +1255,15 @@ 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. 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. + # whole channel. The region selection used to be the one thing left in a + # per-sample loop, because the mode draw was batch-wide everywhere else and + # hoisting it here would have made one decision for the batch. It is per + # sample everywhere now (see `choose_region_modes`), so the loop is gone. + # + # What that costs: the non-finite guard was per sample and is now per channel, + # so one bad sample withholds the whole channel rather than just itself. This + # transform is `max - x`, which is finite whenever its input is, so the guard + # only fires on an input that was already broken. channel = input[:, c] reduce_dims = tuple(range(1, channel.dim())) keep_shape = (-1, *([1] * (channel.dim() - 1))) @@ -1368,14 +1289,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ new_std = inverted.std(dim=reduce_dims).view(keep_shape) inverted = (inverted - new_mean) / (new_std + eps) * orig_stds + orig_means - for i in range(input.shape[0]): - orig = channel[i] - x = inverted[i] - - checked = _select_and_check(self, orig, x, None if seg_mask is None else seg_mask[i]) - if checked is None: - continue - input[i, c] = checked + checked = _select_and_check(self, channel, inverted, seg_mask) + if checked is None: + continue + input[:, c] = checked return input @@ -1413,7 +1330,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1423,9 +1344,7 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency @@ -1574,7 +1493,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, same_on_batch: bool = False, @@ -1595,9 +1518,7 @@ def __init__( self.apply_to_channel = apply_to_channel self.invert = invert self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) def _num_coeffs(self, dim: int) -> int: @@ -1795,7 +1716,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1804,9 +1729,7 @@ def __init__( self.max_clamp_amount = max_clamp_amount self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @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: @@ -1865,17 +1788,17 @@ def __init__( keepdim: bool = True, in_seg: float = 0.0, out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, ) -> None: super().__init__(p=p, p_batch=p_batch, same_on_batch=False, keepdim=keepdim) self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - # Not a constructor parameter: this transform has no mix_in_out knob and the - # registry derives a config's accepted keys from __init__, so adding one there - # would invent a setting. Set here so the region-selection helper's contract holds. - self.mix_in_out = False + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() def apply_transform( @@ -1909,8 +1832,10 @@ def apply_transform( flat = channel.reshape(channel.shape[0], -1) spread = (flat.amax(dim=1) - flat.amin(dim=1)).view(-1, *([1] * (channel.dim() - 1))) channel = torch.where(spread == 0, channel, (channel - mean) / std) - # No mix_in_out here: z-scoring is applied whole, never to a random subset - # of the mask channels. + # `seg_unit` and `bg_split` are accepted here like everywhere else, but a + # z-score applied to a random subset of the units is a strange thing to ask + # for: the point of this transform is to put the whole channel on one scale. + # Left configurable rather than refused, since the surface is uniform. checked = _select_and_check(self, orig, channel, seg_mask) if checked is None: continue diff --git a/smauglab/transforms/gpu/domain_transfer.py b/smauglab/transforms/gpu/domain_transfer.py index 5f972aa..b00a0c0 100644 --- a/smauglab/transforms/gpu/domain_transfer.py +++ b/smauglab/transforms/gpu/domain_transfer.py @@ -36,6 +36,7 @@ import math import os +from collections.abc import Sequence from typing import Any import numpy as np @@ -48,6 +49,16 @@ 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.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, +) from smauglab.transforms.kernels import gaussian_blur3d, random_bias_field3d # The transfer LUT bank is a multi-hundred-MB artefact built offline by @@ -146,6 +157,13 @@ def __init__( p_spatial_mix: float = 0.0, spatial_mix_scale: float = 0.03, spatial_mix_gain: float = 3.0, + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, same_on_batch: bool = False, @@ -223,6 +241,7 @@ def __init__( # 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_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> tuple[Tensor, bool]: @@ -375,5 +394,14 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # transferred distribution SHAPE while placing it back in the z-score scale the # network trains on, so the augmented patch matches the rest of the batch. acc = (acc - acc.mean()) / acc.std().clamp_min(1e-6) * sd + mu - out[b, ch] = acc + + # Region selection after the z-restore above, not after the LUT-space blend: + # at that point `acc` is still in the LUT's [0, 1] domain while `x` is in the + # input's, and a composite has to have both sides on one scale. By here they + # do, which is also why this is the one of the four synthesising transforms + # that needs no `normalize`. + checked = _select_and_check(self, x, acc, seg[b]) + if checked is None: + continue + out[b, ch] = checked return out diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index aac4140..7ffc3c1 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -10,6 +10,21 @@ 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.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, + draw_feather, + feather_alpha, + keep_units, + region_units, + unit_spec, +) from smauglab.transforms.kernels import gaussian_blur3d from smauglab.transforms.rng import shared_choice @@ -249,6 +264,24 @@ class RandomRedistributeSegGPU(ImageOnlyTransform): Mirrors the CPU `RedistributeTransform` behavior using GPU-friendly ops. Works with inputs shaped [N, C, H, W] or [N, C, D, H, W]. + `confine_prob` is the per-sample probability that each region's Gaussian perturbation + is restricted to that region's own voxels rather than added across the whole patch. + It was called `in_seg`, which it should never have been: everywhere else in the + package `in_seg` is the probability of applying the *whole transform* inside the + segmentation, and the two are different questions. The rename is not cosmetic -- + twenty-eight shipped configs set the old name, and had this class also accepted the + standard `in_seg` they would all have kept validating while silently changing meaning, + with no diagnostic possible because the key would be accepted. So this is the one + transform that does **not** take `in_seg` / `out_seg`, and a config still naming + `in_seg` here fails as an unknown parameter and is told what to write instead. + + What it does take is the rest of the region surface, because its work is already + per region: `seg_unit` chooses what a region *is* (one per label value, which is what + it has always used and therefore the default here, or one per connected component), + `bg_split="em"` splits the background into intensity clusters so there is more than + one region of air to redistribute between, `seg_unit_prob` is the chance each region + contributes, and `seg_feather` softens the edge of a confined contribution. + `retain_stats` defaults to True, and wants to stay that way. The whole method operates in a per-sample [0, 1] min-max space and adds a perturbation of up to 2.0 *in that space* -- twice the input's full dynamic range. With @@ -270,7 +303,7 @@ class RandomRedistributeSegGPU(ImageOnlyTransform): def __init__( self, - in_seg: float = 0.2, + confine_prob: float = 0.2, apply_to_channel: Sequence[int] = (0,), retain_stats: bool = True, same_on_batch: bool = False, @@ -279,15 +312,28 @@ def __init__( keepdim: bool = True, std_noise_range: Sequence[float] = (0.1, 0.3), dilation_iterations_range: Sequence[int] = (1, 3), + # `seg_unit` defaults to "label", not the base class's "whole": one region per + # distinct label value is the region set this transform has always used, so this + # is what keeps the default behaviour unchanged. `seg_unit_prob` defaults to 1.0 + # for the same reason -- every region contributing is what it has always done, and + # at 1.0 the keep draw spends nothing, so the default is bit-compatible. + seg_unit: str = "label", + seg_unit_prob: float = 1.0, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, 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 + self.confine_prob = confine_prob self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats self.std_noise_range = std_noise_range self.dilation_iterations_range = dilation_iterations_range + # `in_seg` and `out_seg` are passed as zero and are deliberately not constructor + # parameters; see the class docstring. The rest of the region surface is read. + configure_regions(self, 0.0, 0.0, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() @@ -313,10 +359,34 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ N = input.shape[0] - # One region per label, whichever layout the caller used. Reading regions off - # the channel axis alone collapses a single-channel label map -- which is what - # nnU-Net's trainer passes -- to one foreground blob. - regions = seg_region_masks(seg) + if not len(self.apply_to_channel): + return input + + # The regions to redistribute between, from `gpu/regions.py` rather than from + # `seg_region_masks`. At this class's defaults -- `seg_unit="label"` -- the two + # produce the same set: one region per distinct label value, background included, + # numbered batch-wide. What the decomposition adds is the other two settings: + # `seg_unit="component"` splits a label into its connected pieces, and + # `bg_split="em"` splits the background into intensity clusters, so the + # redistribution has more than one region of air to work with. + # + # The reference image for `bg_split="em"` is the first selected channel. The + # decomposition is shared by every channel of the call -- and by every other + # transform in the pass -- so it cannot be per channel. + units = region_units(seg, input[:, self.apply_to_channel[0] : self.apply_to_channel[0] + 1], unit_spec(self)) + # Only the units some voxel actually lands in. `RegionUnits` always carries a + # background unit, occupied or not, and an unoccupied region is not harmless here: + # `_region_moments` clamps its count to 1 and hands back mean 0, so it would + # contribute a full-amplitude Gaussian at zero. Dropping it is also what makes the + # default region set identical to the `seg_region_masks` one this replaced, whose + # count was the number of label values *present*. + present = torch.bincount(units.ids.reshape(-1), minlength=units.n_units).nonzero().flatten() + region_index = present.view(1, present.numel(), *([1] * spatial_dims)) + regions = units.ids == region_index + # Which regions contribute at all. Spends nothing at the default `seg_unit_prob` + # of 1.0, where every region does. + keep_rows = keep_units(self, units, N, input.device)[:, present].tolist() + confine_sigma = draw_feather(self) # Apply per selected image channel and per batch sample for c in self.apply_to_channel: @@ -344,7 +414,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # sample's draw independent but not in the loop's interleaved order, so a # seeded run differs from the previous release's. has_foreground = (seg.reshape(N, -1) > 0).any(dim=1).tolist() - in_seg_draws = (torch.rand(N, device=input.device) <= self.in_seg).tolist() + confine_draws = (torch.rand(N, device=input.device) <= self.confine_prob).tolist() dilation_draws = torch.randint( self.dilation_iterations_range[0], self.dilation_iterations_range[1] + 1, (N,), device=input.device ).tolist() @@ -368,7 +438,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ continue # Decide redistribution mode once per sample - in_seg_bool = in_seg_draws[b] + confine = confine_draws[b] # Binary masks for regions masks = regions[b] # (R, ...) @@ -415,10 +485,22 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # `counts` is `clamp_min(1)`. to_add = torch.zeros_like(x) rand_sign = 2 * torch.rand(R, device=input.device) - 1 # random sign factor per region + # Feathered only where it is *multiplied in*, never where the statistics + # come from: `_region_moments` relies on the mask being binary (it squares + # it and calls the result a voxel count), so a soft mask there would + # quietly corrupt every per-region mean and standard deviation. Here it + # softens the edge of the confined contribution, which is the whole point. + confine_flat = mask_flat.to(x.dtype) + if confine and confine_sigma > 0.0: + confine_flat = feather_alpha(confine_flat.view(1, R, *x.shape), confine_sigma).view(R, -1) for r in range(R): + # The sign is drawn for every region whether or not it contributes, so + # `seg_unit_prob` does not shift the stream for the regions that do. + if not keep_rows[b][r]: + continue contribution = _normal_pdf(x, means[r], redist_std[r]) * rand_sign[r] - if in_seg_bool: - contribution = contribution * mask_flat[r].view(x.shape) + if confine: + contribution = contribution * confine_flat[r].view(x.shape) to_add += contribution # Normalize to_add if non-zero. `torch.where`, not `if`: the test is one @@ -484,6 +566,16 @@ 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). + in_seg: per-sample probability of confining the synthesis to the foreground labels. + out_seg: per-sample probability of confining it to the background. Both firing means a + mixture of the two; neither means the whole patch, which is the default and what + this transform has always done. + seg_unit: what a unit is when only part of a side is used -- "whole", "label" or + "component". `seg_unit_prob` is the chance each unit takes part. + bg_split: "em" splits the background into intensity clusters so `out_seg` can land on + part of it. Not to be confused with the generator's own `em_*` parameters. + seg_feather: Gaussian sigma in voxels for the boundary between the composited region + and the untouched one. 0 is a hard edge. 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 @@ -505,6 +597,13 @@ def __init__( label_remap_prob: float = 0.5, min_label_voxels: int = 4, label_classes: list[int] | None = None, + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -527,6 +626,7 @@ def __init__( self.label_remap_prob = label_remap_prob self.min_label_voxels = min_label_voxels self.label_classes = label_classes + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() @@ -680,7 +780,17 @@ def apply_transform( synth_z = ((synth - b_mean) / b_std * flat_m_all).reshape(B, 1, D, H, W) out = input.clone() - out[:, 0:1] = synth_z.to(input.dtype) + # Region selection last, after the foreground z-score above -- which is this + # transform's stats restore -- as it is in every other transform. `normalize=True` + # because the synthesis is masked to zero on the background while the input's + # background is z-scored air: compositing them raw would put an intensity step on + # the mask boundary that has nothing to do with the augmentation. + # + # The non-finite guard is new here. This used to write `synth_z` unconditionally, + # so a NaN out of the synthesis reached the batch. + checked = _select_and_check(self, input[:, 0], synth_z[:, 0], seg_raw, normalize=True) + if checked is not None: + out[:, 0:1] = checked.unsqueeze(1).to(input.dtype) return out diff --git a/smauglab/transforms/gpu/mixing.py b/smauglab/transforms/gpu/mixing.py index c5fa185..e79d6df 100644 --- a/smauglab/transforms/gpu/mixing.py +++ b/smauglab/transforms/gpu/mixing.py @@ -83,8 +83,23 @@ class Mixing(Protocol): _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]``. +def _bound_text(value: float | None) -> str: + """A bound as it appears in a message: ``0``, ``1``, ``inf``. + + `%g` rather than `str`, so the default bounds still render as the ``[0, 1]`` the + error message has always said rather than as ``[0.0, 1.0]``. + """ + return "inf" if value is None else f"{value:g}" + + +def normalise_blend_range( + spec: MixAlphaSpec, + key: str = "mix_alpha", + owner: str = "", + hint: str = _MIX_ALPHA_HINT, + bounds: tuple[float, float | None] = (0.0, 1.0), +) -> tuple[float, float]: + """A scalar or ``[lo, hi]`` as a validated pair inside `bounds`. 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 @@ -93,6 +108,11 @@ def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str `key` and `owner` are for the message alone. `AnatomicalLabelOverlay.blend_strength` is a different blend with the same shape of argument, and shares this. + + `bounds` is what makes it shareable by a parameter that is *not* a weight: `gpu/regions.py` + passes ``(0.0, None)`` for `seg_feather`, a Gaussian sigma in voxels with no upper + limit. The default is the unit interval every blend weight here needs, so the three + existing callers are unaffected -- including the message, which still reads ``[0, 1]``. """ where = f"{owner}: " if owner else "" if isinstance(spec, bool): # an int subtype, and never a blend weight anyone meant @@ -111,9 +131,10 @@ def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str 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: + low, high = bounds + if lo < low or (high is not None and hi > high): suffix = f" {hint}" if hint else "" - raise ValueError(f"{where}{key} must lie inside [0, 1], got [{lo}, {hi}].{suffix}") + raise ValueError(f"{where}{key} must lie inside [{_bound_text(low)}, {_bound_text(high)}], got [{lo}, {hi}].{suffix}") return (lo, hi) diff --git a/smauglab/transforms/gpu/palette/transform.py b/smauglab/transforms/gpu/palette/transform.py index da16cdb..5af481a 100644 --- a/smauglab/transforms/gpu/palette/transform.py +++ b/smauglab/transforms/gpu/palette/transform.py @@ -18,6 +18,16 @@ ) from smauglab.transforms.gpu.palette.factory import make_initial, make_overlay, make_refinement from smauglab.transforms.gpu.palette.overlay import AnatomicalLabelOverlay +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, +) from smauglab.transforms.rng import shared_choice #: The composition `RandomPaletteGPU` hard-codes, which is what this defaults to. @@ -87,6 +97,13 @@ 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), + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -103,6 +120,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_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() @@ -180,5 +198,10 @@ def apply_transform( synth_z = ((synth - b_mean) / b_std * flat_m_all).reshape(B, 1, D, H, W) out = input.clone() - out[:, 0:1] = synth_z.to(input.dtype) + # Region selection last, after the foreground z-score -- this transform's stats + # restore -- as everywhere else. `normalize=True` because the synthesis is masked + # to zero on the background while the input's background is z-scored air. + checked = _select_and_check(self, input[:, 0], synth_z[:, 0], seg_raw, normalize=True) + if checked is not None: + out[:, 0:1] = checked.unsqueeze(1).to(input.dtype) return out diff --git a/smauglab/transforms/gpu/regions.py b/smauglab/transforms/gpu/regions.py new file mode 100644 index 0000000..f3a291b --- /dev/null +++ b/smauglab/transforms/gpu/regions.py @@ -0,0 +1,894 @@ +"""Where a transform applies, and how sharply it stops applying. + +Three parameters predate this module -- `in_seg`, `out_seg` and `mix_in_out` -- and +between them they expressed about a third of what they looked like they expressed: + +* `in_seg` and `out_seg` were two independent uniforms, and *both fired* fell through + to the same ``"all"`` as *neither fired*. "A mixture of foreground and background" + was not expressible, and `in_seg=1, out_seg=1` was indistinguishable from + `in_seg=0, out_seg=0`. +* `mix_in_out` keyed its keep/drop draw on `seg_mask.shape[:2]`, the **channel** axis. + For the single-channel integer label map nnU-Net's trainer passes, that is `(B, 1)` + -- one draw per sample, deleting the whole mask half the time. So ``"in"`` applied + the transform nowhere and ``"out"`` applied it everywhere, and the thirteen shipped + configs that set `mix_in_out: true` were running a coin flip rather than a label mix + in every nnU-Net run. It only ever meant "per label" on a one-hot mask. +* The background was one blob. `out_seg` took all of it or none of it. +* The composite was `m * transformed + (1 - m) * orig` with a **binary** `m`, so every + region-selected augmentation left a one-voxel discontinuity sitting exactly on the + label boundary -- an edge perfectly correlated with the segmentation target, which is + the one thing an augmentation must not hand the network. + +The replacement separates four questions that were tangled together: + +================ ============================================================ +`in_seg`, **where**: foreground labels, background, a mixture of both, +`out_seg` or the whole patch. Both fired now means ``"both"``. +`seg_unit`, **at what granularity**: the whole side as one piece, one unit +`seg_unit_prob` per label value, or one unit per connected component. +`bg_split`, **how finely the background is divided**: one blob, or EM +`bg_em_clusters` intensity clusters, SynthSeg-style (Billot et al. 2023, 5.4). +`seg_feather` **how sharp the boundary is**: a hard step, or a Gaussian ramp. +================ ============================================================ + +Everything resolves to a single integer **unit id map** (:class:`RegionUnits`) and one +gather off it. Nothing materialises per-unit volumes: a noisy mask under +``seg_unit="component"`` can hold hundreds of components, and `[2, 300, 128, 128, 128]` +of bool is 1.2 GB. + +Three properties this module goes out of its way to keep, each of which the code it +replaces broke: + +**A knob that is off spends no random numbers.** `_choose_region_mode` drew its two +uniforms whenever a mask was attached, including at `in_seg = out_seg = 0`, so an +unselected run and a selected one at the same seed were never comparable. That is the +defect `gpu/mixing.py` removed for `mix_prob`, and :func:`choose_region_modes` returns +before drawing anything for the same reason. + +**The number of uniforms a draw spends is a function of the configuration, never of the +data.** A per-unit keep draw would spend one uniform per label *present*, so a patch +with four labels and one with five would leave the stream in different places and every +transform after this one would diverge on the data. :data:`UNIT_DRAW_WIDTH` fixes the +width per decomposition kind instead. + +**The decomposition is computed once per pipeline pass, not once per transform.** See +:func:`region_units`. Connected components and EM each cost tens of milliseconds on a +128^3 patch; a pipeline runs twenty-odd region-selecting transforms over one batch. + +One consequence worth stating plainly rather than leaving to be discovered: +``bg_split="em"`` is the only decomposition that reads the image, and the image moves as +the pass runs -- every transform's first statement is `input = input.clone()`. The +clusters are therefore fitted **once per pass, from the image as it stood when the first +region-selecting transform asked**, and not re-fitted afterwards. That is deliberate +twice over. Re-fitting is twenty-odd EM runs a batch, each with a 100 000-element +subsample draw and a host synchronisation per sample; and a single fit means two +region-selected transforms in the same pass cut the background along the same lines +instead of each along its own. What it costs is that a transform late in the pass +clusters an image a few transforms old. The units are a partition to switch on and off, +not a measurement, so that is the cheap half of the trade. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, NamedTuple, Protocol, Union + +import torch +from torch import Tensor + +from smauglab.transforms.gpu.mixing import normalise_blend_range +from smauglab.transforms.kernels import gaussian_blur3d +from smauglab.transforms.rng import next_shared_seed + +#: A fixed sigma in voxels, or a ``[lo, hi]`` range drawn from per call. Mirrors +#: `MixAlphaSpec` in gpu/mixing.py, which is the same shape of argument for a weight +#: rather than a width. +SegFeatherSpec = Union[float, Sequence[float]] + +SEG_UNITS = ("whole", "label", "component") +BG_SPLITS = ("none", "em") + +DEFAULT_SEG_UNIT = "whole" +DEFAULT_SEG_UNIT_PROB = 0.5 +DEFAULT_BG_SPLIT = "none" +DEFAULT_BG_EM_CLUSTERS = (3, 10) +DEFAULT_SEG_FEATHER = 0.0 + +#: How many keep/drop uniforms a draw spends, per decomposition kind. Fixed per *kind* +#: so the count is a function of the configuration and never of the segmentation -- see +#: the module docstring. A unit id at or above the width folds into the last unit, which +#: only bites a patch with more than 32 label values or more than 256 components per +#: side, and then only by tying their fates together. +UNIT_DRAW_WIDTH = {"whole": 2, "label": 32, "component": 256} + +#: Mode codes, as they travel in the `[N]` tensor `choose_region_modes` returns. +MODE_ALL = 0 +MODE_IN = 1 +MODE_OUT = 2 +MODE_BOTH = 3 + +#: A label value no real label can take, for the gated shifts in +#: :func:`connected_components_3d`: the out-of-volume neighbour must never compare equal +#: to the voxel it is shifted past. +_NO_LABEL = -1 + +#: Where the attribute cache lives on a mask tensor. See :func:`region_units`. +_CACHE_ATTR = "_smauglab_region_units" + +#: Appended to a `seg_feather` bound error, because "sigma in voxels" is not what the +#: name suggests to someone who has just read `mix_alpha`. +_FEATHER_HINT = "It is a Gaussian sigma in voxels, not a weight: 0 is a hard boundary, and there is no upper limit." + + +class RegionUnits(NamedTuple): + """A segmentation decomposed into the units a mode can switch on and off. + + `ids` is dense in ``[0, n_units)`` and background units come **first**, so "is this + unit foreground" is the scalar comparison ``id >= n_background`` rather than a + per-unit tensor. Label units are numbered batch-wide, so a given label value is the + same unit in every sample of a batch -- the choice `seg_region_masks` already + documents for its channel ordering. Component and EM-cluster units cannot be: + they are per-sample by nature, and sample `b`'s unit `k` is unrelated to sample + `b+1`'s. The keep draw is per sample either way, so that costs nothing. + """ + + ids: Tensor + n_units: int + n_background: int + + +class UnitSpec(NamedTuple): + """What a decomposition depends on, and therefore the cache key. + + `seg_unit_prob`, `in_seg`, `out_seg` and `seg_feather` are deliberately absent: they + are redrawn per call and must not fragment the cache. + """ + + seg_unit: str + bg_split: str + bg_em_clusters: tuple[int, int] + + +class RegionSelecting(Protocol): + """What the helpers below need off the transform they are handed. + + A Protocol rather than `ImageOnlyTransform`, for the same reason `mixing.Mixing` 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 neither of those is a float. Naming them here is both what + makes the arithmetic below type-check and a statement of the requirement. + """ + + in_seg: float + out_seg: float + seg_unit: str + seg_unit_prob: float + bg_split: str + bg_em_clusters: tuple[int, int] + seg_feather: tuple[float, float] + + +# ── validation ─────────────────────────────────────────────────────────────── + + +def _one_of(value: Any, choices: tuple[str, ...], key: str, owner: str) -> str: + where = f"{owner}: " if owner else "" + if not isinstance(value, str) or value not in choices: + raise ValueError(f"{where}{key} must be one of {', '.join(choices)}, got {value!r}") + return value + + +def _cluster_range(spec: Any, owner: str) -> tuple[int, int]: + """A ``[lo, hi]`` pair of positive cluster counts. + + Not `normalise_blend_range`: these are counts, and silently accepting a float would + put a non-integer into `torch.randint`'s bounds. + """ + where = f"{owner}: " if owner else "" + if isinstance(spec, (bool, int, float)): + raise ValueError(f"{where}bg_em_clusters must be a [lo, hi] pair of counts, got {spec!r}") # noqa: TRY004 + try: + values = [int(v) for v in spec] + except (TypeError, ValueError) as exc: + raise ValueError(f"{where}bg_em_clusters must be a [lo, hi] pair of counts, got {spec!r}") from exc + if len(values) != 2: + raise ValueError(f"{where}bg_em_clusters must be a [lo, hi] pair of counts, got {len(values)} values: {spec!r}") + lo, hi = values + if lo < 1 or hi < lo: + raise ValueError(f"{where}bg_em_clusters must be a non-decreasing pair of counts >= 1, got [{lo}, {hi}]") + if hi > UNIT_DRAW_WIDTH["whole"] + max(UNIT_DRAW_WIDTH.values()): + # Only a sanity bound: a background split into thousands of clusters is a + # configuration mistake, and the keep draw would fold most of them together. + raise ValueError(f"{where}bg_em_clusters upper bound {hi} is implausibly large") + return (lo, hi) + + +def configure_regions( + transform: RegionSelecting, + in_seg: float, + out_seg: float, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, +) -> None: + """Validate and assign the region parameters. One line per constructor. + + Every value check in this package happens in a constructor -- `smauglab/config.py` + only ever asks whether a key is *accepted*, never what its value is. A pipeline is + built before the first batch, so a bad value still fails long before training. + + `in_seg` and `out_seg` used to be clamped into range silently + (``max(0.0, min(1.0, p))``), which quietly accepted `in_seg: 5.0` and + `in_seg: true`. They are rejected now. + """ + owner = type(transform).__name__ + transform.in_seg = normalise_blend_range(in_seg, key="in_seg", owner=owner, hint="")[0] + transform.out_seg = normalise_blend_range(out_seg, key="out_seg", owner=owner, hint="")[0] + transform.seg_unit = _one_of(seg_unit, SEG_UNITS, "seg_unit", owner) + transform.seg_unit_prob = normalise_blend_range(seg_unit_prob, key="seg_unit_prob", owner=owner, hint="")[0] + transform.bg_split = _one_of(bg_split, BG_SPLITS, "bg_split", owner) + transform.bg_em_clusters = _cluster_range(bg_em_clusters, owner) + transform.seg_feather = normalise_blend_range(seg_feather, key="seg_feather", owner=owner, hint=_FEATHER_HINT, bounds=(0.0, None)) + + +def unit_spec(transform: RegionSelecting) -> UnitSpec: + """The decomposition `transform` asks for, which is also its cache key.""" + return UnitSpec(transform.seg_unit, transform.bg_split, transform.bg_em_clusters) + + +def region_selection_active(transform: RegionSelecting) -> bool: + """Whether either probability can fire. False means every helper here is a no-op.""" + return transform.in_seg > 0.0 or transform.out_seg > 0.0 + + +# ── segmentation layout ────────────────────────────────────────────────────── + + +def label_map_from(mask: Tensor) -> Tensor: + """A `[B, 1, *spatial]` integer label map, background 0, from any of three layouts. + + The layouts reach a transform interchangeably and the values alone cannot always + tell them apart: + + * ``[B, 1, *spatial]`` -- an ordinary label map. The ids are the values and + background is 0. Rounded first: a mask that has been through a geometric transform + comes back as float, and unrounded floats would invent labels. + * ``[B, C, *spatial]`` that leaves some voxel unset -- a background-*implicit* + one-hot, this repository's convention (`collapse_onehot_to_index` in + gpu/fromSeg.py): channel `c` encodes label `c + 1`. + * ``[B, C, *spatial]`` that leaves no voxel unset -- a background-*inclusive* + one-hot, which is what `seg_region_masks` emits for a single-channel label map and + what `unit_tests/test_seg_layout.py` builds. Channel 0 **is** background. + + The third is the one `_foreground` already had to detect by coverage rather than by + value, and the one `collapse_onehot_to_index` still gets wrong: it does `argmax + 1` + unconditionally, so a background-inclusive one-hot comes back with every label + shifted by one and background relabelled 1. That is why `em_subdivide_labels` carries + a fallback for "the real background is no longer 0". Deriving the map here, once, is + what lets everything else in this module be a read off it. + + One case no function can read off the values: a background-implicit one-hot patch + labelled edge to edge is indistinguishable from a background-inclusive one. A real + segmentation patch is not labelled edge to edge, which is the same admission + `_foreground`'s docstring has always carried. + """ + if mask.dim() < 3: + raise ValueError(f"expected [B, C, *spatial], got {tuple(mask.shape)}") + if mask.shape[1] == 1: + return mask.round().long().clamp_min(0) + + covered = mask.amax(dim=1, keepdim=True) > 0 + index = mask.argmax(dim=1, keepdim=True).long() + # `torch.where`, not `if bool(covered.all())`: reading that test back to Python + # stalls the host on everything queued behind it, and both branches are a cheap + # elementwise pass. Same reason `_foreground` is written this way. + return torch.where(covered.all(), index, (index + 1) * covered) + + +def _foreground(mask: Tensor, dim: int) -> Tensor: + """Which voxels the segmentation covers, reducing over the class axis. + + This used to be `torch.argmax(mask, dim) > 0`, which is only "is anything labelled + here" if class 0 is background -- and it is not: + + * For a single-channel `[B, 1, D, H, W]` mask (an ordinary nnU-Net target, and what + the tests build) `argmax` over a length-1 axis is always 0, so the result was + **all False**. `in_seg` then applied the transform nowhere and `out_seg` applied + it everywhere: the two knobs did nothing and the opposite of nothing. + * For a one-hot mask, this repository's convention (`collapse_onehot_to_index` in + gpu/fromSeg.py) is that channel `c` encodes label `c + 1` with background + implicit, so `argmax == 0` is a real foreground class and was being dropped. + + `amax > 0` asks the question that was meant, and matches what + `collapse_onehot_to_index` already does with `seg_raw.any(dim=1)`. + + That still leaves one layout it cannot read off the values alone: a one-hot + that *includes* a background channel, which is what `seg_region_masks` emits + for a single-channel label map and what `unit_tests/test_seg_layout.py` + builds. Every voxel then has some channel set, `amax > 0` is True everywhere, + and `in_seg` applies the transform to the whole patch while `out_seg` applies + it nowhere -- the exact failure this function was written to fix, in the one + layout it was not checked against. + + A background-inclusive one-hot is recognisable: it is multi-channel and leaves + no voxel unset, because the background channel covers whatever the foreground + channels do not. A background-implicit one-hot always has all-zero background + voxels unless the patch is labelled edge to edge, which a real segmentation + patch is not. So full coverage plus more than one channel means channel 0 is + background, and the reduction skips it. + + Kept as its own function, rather than folded into :func:`label_map_from`, because + `dim` is not always 1: `_apply_region_mode`'s unbatched path reduces a `[C, *spatial]` + mask over axis 0. The two agree on every layout, which + `unit_tests/test_foreground_layouts.py` is the check on. + """ + covered = mask.amax(dim=dim) > 0 + if mask.shape[dim] == 1: + return covered + # `torch.where`, not `if bool(...)`: reading that test back to Python stalls the host + # on everything queued behind it. Both branches are a reduction over the class axis, + # so computing the unused one costs a fraction of what the stall does. + without_background = mask.narrow(dim, 1, mask.shape[dim] - 1).amax(dim=dim) > 0 + return torch.where(covered.all(), without_background, covered) + + +def densify_region_ids(region_ids: Tensor) -> tuple[Tensor, int]: + """Remap sparse non-negative integer region ids onto a contiguous [0, R) range. + + `em_subdivide_labels` encodes a fine id as `parent_idx * mult + assign`, which + is sparse. Anything that indexes a per-region tensor of length R with these -- the + keep/drop gather here, `signed_alpha_affine_remap` in the PALETTE blocks -- would + index out of bounds on a sparse id. + + Lives here rather than in `gpu/palette/partitioners.py`, where it was written, + because this module cannot import that one: `palette/base.py` reaches `gpu/base.py`, + which imports this module for its "off" defaults. Re-exported from there. + """ + # A lookup table off a bincount, not `torch.unique(return_inverse=True)`: the ids + # are small non-negative integers, so counting the bins and numbering the + # non-empty ones gives the same ascending remap without radix-sorting two + # million voxels. + counts = torch.bincount(region_ids.reshape(-1)) + present = counts.nonzero().flatten() + lookup = torch.zeros(counts.numel(), dtype=torch.long, device=region_ids.device) + lookup[present] = torch.arange(present.numel(), device=region_ids.device) + return lookup[region_ids], int(present.numel()) + + +# ── connected components ───────────────────────────────────────────────────── + +#: Hard ceiling on the propagation below. Pointer jumping squares the reach each round, +#: so even a serpentine component spanning a 128^3 patch settles in about twenty; the cap +#: only exists so a bug cannot hang a training run. +_CC_MAX_ITERS = 64 + +#: How often the fixpoint test runs. `bool(...)` on a device tensor stalls the host until +#: the queue drains, and the loop body is six shifts over the volume, so testing every +#: round would spend a synchronisation to save work that is already cheap. Compact blobs +#: converge in a handful of rounds, which is why the test exists at all rather than +#: running a fixed count. Same trade `_kmeans_1d` made in the other direction. +_CC_CHECK_EVERY = 4 + + +def _shift(volume: Tensor, dim: int, delta: int, fill: Any) -> Tensor: + """`volume` shifted along `dim` so that index `i` holds what index `i - delta` held. + + `fill` enters at the face that is uncovered. Not `torch.roll`, which wraps the far + face round to the near one and would join two components that only touch through + opposite sides of the patch. + """ + size = volume.shape[dim] + step = abs(delta) + kept = volume.narrow(dim, 0, size - step) if delta > 0 else volume.narrow(dim, step, size - step) + pad_shape = list(volume.shape) + pad_shape[dim] = step + block = torch.full(pad_shape, fill, dtype=volume.dtype, device=volume.device) + return torch.cat([block, kept] if delta > 0 else [kept, block], dim=dim) + + +def connected_components_3d(labels: Tensor, foreground: Tensor) -> tuple[Tensor, int]: + """Connected components of each label value, 6-connected, per sample. + + `labels` is `[B, 1, D, H, W]` integer and `foreground` the bool mask of the voxels to + partition. Returns `(ids, n_max)` where `ids` is `[B, 1, D, H, W]`, numbered from 0 + within each sample and `-1` outside `foreground`, and `n_max` is the largest + per-sample count -- which is what sizes the keep/drop draw. + + **6-connected**, matching `scipy.ndimage.label`'s default structuring element + (`generate_binary_structure(3, 1)`), which is also the conventional reading of + "connected components of a label": two voxels touching only at a corner are two + components. The obvious implementation -- one `max_pool3d(kernel=3)` -- is a 3^3 + neighbourhood and therefore 26-connected, which merges them. Six explicit shifts + cost twice the kernels of three axis pools and buy the right answer. + + Components respect label boundaries without a pass per label value. The gate is the + `same` mask below: a neighbour contributes only where it carries the *same* label, so + one propagation partitions every label at once. Doing it with pooled values instead + would need a pass per label (a spine patch has five to twenty-five), because a pool + reports the largest value in the neighbourhood without saying which voxel it came + from, and a higher-labelled neighbour would mask the same-labelled one. + + The algorithm is label-equivalence propagation: every voxel starts as its own + component, named by its flat index; each round takes the largest name reachable from + a same-label neighbour, then **jumps to that name's own current name**, which doubles + the distance information travels per round and turns a linear walk along a long thin + structure into a logarithmic one. + + Indices are `int64` throughout. The float alternative is tempting because + `max_pool3d` requires it, but `float32` represents integers exactly only to 2^24, and + a 256^3 patch has exactly that many voxels -- so a larger patch would silently merge + components whose flat indices round together. + """ + batch = labels.shape[0] + spatial = labels.shape[2:] + voxels = int(labels[0, 0].numel()) + device = labels.device + + # Gated once, outside the loop: the label map does not move, so the six + # "neighbour carries the same label, and both of us are foreground" masks are + # constants of the call. Recomputing them per round would double its cost. + same = [] + for dim in (2, 3, 4): + for delta in (1, -1): + neighbour_label = _shift(labels, dim, delta, _NO_LABEL) + neighbour_fg = _shift(foreground, dim, delta, False) + same.append(((neighbour_label == labels) & neighbour_fg & foreground, dim, delta)) + + # Globally unique seeds, not per-sample ones: the densify below counts over the whole + # batch, so two samples whose components happened to root at the same within-sample + # index would be counted as one. + ids = torch.arange(batch * voxels, device=device, dtype=torch.long).view(batch, 1, *spatial) + ids = torch.where(foreground, ids, torch.full_like(ids, -1)) + + for iteration in range(1, _CC_MAX_ITERS + 1): + best = ids + for gate, dim, delta in same: + best = torch.maximum(best, torch.where(gate, _shift(ids, dim, delta, -1), torch.full_like(ids, -1))) + flat = best.flatten(1) + # The name that our current name currently answers to. `clamp_min(0)` only keeps + # the gather in range for the background entries, which the `where` then drops. + jumped = torch.gather(flat, 1, (flat - batch_offset(batch, voxels, device)).clamp_min(0)) + advanced = torch.maximum(flat, torch.where(flat >= 0, jumped, flat)).view_as(ids) + advanced = torch.where(foreground, advanced, torch.full_like(advanced, -1)) + settled = iteration % _CC_CHECK_EVERY == 0 and bool(torch.equal(advanced, ids)) + ids = advanced + if settled: + break + + # Densified over the batch, then rebased per sample. Global ids are ascending in the + # root's flat index and every index of sample b is below every index of sample b+1, so + # each sample's ids are a contiguous ascending block and subtracting its lowest gives + # a 0-based numbering without a second densify per sample. + dense, _total = densify_region_ids(ids.clamp_min(0)) + dense = torch.where(foreground, dense, torch.full_like(dense, -1)) + flat_dense = dense.flatten(1) + sentinel = torch.full_like(flat_dense, torch.iinfo(torch.int64).max) + lowest = torch.where(flat_dense >= 0, flat_dense, sentinel).min(dim=1).values + highest = flat_dense.max(dim=1).values + counts = torch.where(highest >= lowest, highest - lowest + 1, torch.zeros_like(highest)) + rebased = torch.where(foreground, dense - lowest.view(-1, *([1] * (dense.dim() - 1))), torch.full_like(dense, -1)) + return rebased, max(int(counts.max()), 1) + + +def batch_offset(batch: int, voxels: int, device: torch.device) -> Tensor: + """`[B, 1]` of each sample's first global seed index. + + The names in :func:`connected_components_3d` are global so the densify can count them + together, but `torch.gather` indexes within a row, so the row's offset comes off again + before the lookup. + """ + return (torch.arange(batch, device=device, dtype=torch.long) * voxels).view(batch, 1) + + +# ── background decomposition ───────────────────────────────────────────────── + + +def background_em_units(image: Tensor, labels: Tensor, clusters_range: tuple[int, int]) -> tuple[Tensor, int]: + """The background split into intensity-coherent subregions by EM, SynthSeg-style. + + Returns `(ids, n)` with `ids` numbered from 0 on background voxels and 0 elsewhere + (foreground ids are never read: the caller overwrites them with the foreground + numbering). + + `em_subdivide_labels` does the work -- this is a two-label call that asks it to split + label 0 into `clusters_range` clusters and leave the foreground as a single one, which + is exactly the paper's background handling without its foreground handling. + + Two things are wrapped around it rather than left to it: + + * **The draw is forked.** It spends one `randint` per sample *and* a + `max_fit_voxels`-element subsample index -- a hundred thousand uniforms, which the + background of a 128^3 patch always triggers. Since the result is cached, a miss + would spend them and a hit would not, which would make the global RNG stream a + function of cache state: of which transform asked first, of which rows `p` selected, + of whether a geometric transform replaced the mask. Forking keeps the stream exactly + where it was found. `next_shared_seed` rather than an arbitrary seed so DDP ranks + still agree. + * **A patch with no background is left alone.** When the configured background label + is absent, `em_subdivide_labels` falls back to the largest-area label, which here + would silently split a *foreground* structure into three to ten pieces and call them + background. + + `same_on_batch` is fixed at False: the decomposition is shared by every transform in + the pass, so it cannot follow any one transform's flag. + """ + # Imported here, not at module scope. `smauglab/transforms/synthseg/__init__.py` + # imports synthseg/transforms.py, which imports gpu/base.py -- and gpu/base.py imports + # this module for its "off" defaults. A module-scope import would close that loop and + # break every `import smauglab.transforms.gpu.base`; `unit_tests/test_imports.py` is + # the canary. By the time any transform calls this, the package is fully imported. + from smauglab.transforms.synthseg.functional import em_subdivide_labels + + background = labels <= 0 + if not bool(background.any()): + return torch.zeros_like(labels), 1 + + binary = (~background).long() + device = image.device + with torch.random.fork_rng(devices=[device] if device.type == "cuda" else [], enabled=True): + torch.manual_seed(next_shared_seed()) + fine, _generation, _merge = em_subdivide_labels( + image, + binary, + n_foreground_clusters=1, + background_clusters_range=clusters_range, + background_label=0, + same_on_batch=False, + ) + # Background is parent index 0, so its fine ids are `0 * mult + assign` -- already the + # 0-based cluster numbering this returns. Foreground ids are zeroed because the caller + # replaces them. + ids = torch.where(background, fine, torch.zeros_like(fine)) + return ids, max(int(ids.max()) + 1, 1) + + +# ── the decomposition, and its cache ───────────────────────────────────────── + + +def _label_units(labels: Tensor) -> tuple[Tensor, int]: + """One unit per distinct positive label value, numbered batch-wide from 0.""" + counts = torch.bincount(labels.reshape(-1).clamp_min(0)) + present = counts.nonzero().flatten() + positive = present[present > 0] + lookup = torch.zeros(counts.numel(), dtype=torch.long, device=labels.device) + lookup[positive] = torch.arange(positive.numel(), device=labels.device) + return lookup[labels.clamp_min(0)], max(int(positive.numel()), 1) + + +def _as_batched(volume: Tensor) -> Tensor: + """`[C, *spatial]` or `[B, C, *spatial]` as `[B, C, *spatial]`. + + The region helpers are called from two shapes of `apply_transform`: the vectorised + ones hand over the whole batch, and the per-sample ones (`RandomInverseGPU`, + `RandomDomainTransferGPU`) hand over one sample with no batch axis. + """ + return volume if volume.dim() == 5 else volume.unsqueeze(0) + + +def _build_units(mask: Tensor, image: Tensor | None, spec: UnitSpec) -> RegionUnits: + labels = label_map_from(_as_batched(mask)) + foreground = labels > 0 + + if spec.bg_split == "em" and image is not None: + background_ids, n_background = background_em_units(_as_batched(image), labels, spec.bg_em_clusters) + else: + background_ids, n_background = torch.zeros_like(labels), 1 + + if spec.seg_unit == "label": + foreground_ids, n_foreground = _label_units(labels) + elif spec.seg_unit == "component": + foreground_ids, n_foreground = connected_components_3d(labels, foreground) + else: + foreground_ids, n_foreground = torch.zeros_like(labels), 1 + + ids = torch.where(foreground, n_background + foreground_ids.clamp_min(0), background_ids) + return RegionUnits(ids, n_background + n_foreground, n_background) + + +def region_units(mask: Tensor, image: Tensor | None, spec: UnitSpec) -> RegionUnits: + """The decomposition of `mask`, computed once per mask tensor per spec. + + Connected components and EM each cost tens of milliseconds on a 128^3 patch, and a + pipeline runs twenty-odd region-selecting transforms over one batch. Without a cache + the feature is not trainable; with one, every transform of a pass shares the same + units, which is also what makes two region-selected transforms cut the background + along the same lines rather than each along its own. + + **The cache is an attribute on the mask tensor**, keyed by spec. That choice is load + bearing in three ways: + + * **Lifetime.** It dies with the tensor, so there is nothing to bound, evict or clear. + * **Invalidation.** `MaskSequentialOpsCustom.transform` hands back the *same* tensor + object for an image-only module and a *new* one for a geometric module, so the + identity changes exactly when the mask changes. A geometric transform mid-pipeline + invalidates this for free. + * **Not a dict keyed on tensors.** `Tensor.__hash__` is identity-based but + `Tensor.__eq__` is elementwise, so a hash-bucket collision between two different + masks makes `dict` evaluate `bool(a == b)` and raise *"Boolean value of Tensor with + more than one value is ambiguous"*. Measured: a `WeakKeyDictionary` with five + hundred same-shaped tensor keys raises reliably. An attribute has no buckets. + + Keyed on the mask **as the transform was handed it**, which is what keeps the image + and the units row-aligned. kornia slices the image down to the rows that drew below + `p` and `_seg_for_applied_rows` slices the mask to match, so a decomposition built + from the full batch would pair B-row units with a B'-row image -- and `bg_split="em"` + reads that image. When every row applies, which is the common case and always true at + `p = 1`, the slice *is* the full mask object and every module of the pass hits the + same entry. A module whose row selection is a strict subset pays for its own + decomposition, which is the honest price of having the rows line up. + """ + cached = getattr(mask, _CACHE_ATTR, None) + if cached is not None and spec in cached: + return cached[spec] + units = _build_units(mask, image, spec) + if cached is not None: + cached[spec] = units + else: + try: + setattr(mask, _CACHE_ATTR, {spec: units}) + except AttributeError: # pragma: no cover -- a tensor subclass that forbids it + pass + return units + + +# ── the per-call draws ─────────────────────────────────────────────────────── + + +def choose_region_modes(transform: RegionSelecting, batch: int, device: torch.device) -> Tensor | None: + """One mode per sample, or None when neither probability can fire. + + Returns `MODE_ALL` / `MODE_IN` / `MODE_OUT` / `MODE_BOTH` codes. **Both fired means + `MODE_BOTH`**, where it used to fall through to `MODE_ALL` -- so `in_seg=1, out_seg=1` + stopped being indistinguishable from `in_seg=0, out_seg=0`. + + Returns None, **having drawn nothing**, when both probabilities are zero. The version + this replaces drew its two uniforms whenever a mask was attached, which is the defect + `984f037` removed for `mix_prob`: with the draws unconditional, a run with region + selection off and a run with it on were never on the same stream, so no sweep could be + compared against its own baseline at a fixed seed. + + Per sample, not per batch. Ten of the eleven call sites used to take one decision for + the whole batch, so `in_seg=0.5` meant "half the batches" rather than "half the + samples" -- while `p` and `mix_prob` were per-sample all along. On the image's device + and with no `dtype=`, mirroring `mix_with_original`, so a half-precision batch stays on + the same stream as a single-precision one. + """ + if not region_selection_active(transform): + return None + draws = torch.rand(batch, 2, device=device) + fired_in = draws[:, 0] < transform.in_seg + fired_out = draws[:, 1] < transform.out_seg + # The code arithmetic is the definition: 0 neither, 1 in, 2 out, 3 both. + return fired_in.long() + 2 * fired_out.long() + + +def draw_feather(transform: RegionSelecting) -> float: + """The feather sigma for this call, in voxels. Zero means a hard boundary. + + One draw per call rather than per sample: `gaussian_blur3d` reads its sigmas back to + the host (`.tolist()`), so a per-sample sigma would be one blur call and one + synchronisation per sample. The modes and the per-unit keeps stay per sample; a + boundary *width* is not a decision about where. From the CPU default generator as a + Python float, the way `blend_one` draws its gate, so nothing has to be read off the + device. + + Spends nothing when `seg_feather` is zero, and otherwise spends exactly one uniform + whatever the modes came out as. Gating the draw on "did any sample pick a mode that + uses it" would make the stream a function of the data. Drawn even when `lo == hi`, for + the reason `mix_alpha` is: otherwise the stream becomes a function of the range, and + two runs differing only in `seg_feather` would not be comparable. + """ + low, high = transform.seg_feather + if high <= 0.0: + return 0.0 + return float(torch.rand(1).item()) * (high - low) + low + + +def _keep_draw(transform: RegionSelecting, batch: int, columns: int, width: int, device: torch.device) -> Tensor: + """`[batch, columns]` of which units are kept, at a fixed draw width. + + Spends nothing at `seg_unit_prob >= 1`, which means "every unit" and is therefore the + setting that turns granularity off. Same rule as `in_seg = out_seg = 0` spending + nothing in :func:`choose_region_modes`, and it is what lets `seg_unit="label"` at + probability 1 reproduce `"whole"` bit for bit rather than merely in distribution. + """ + if transform.seg_unit_prob >= 1.0: + return torch.ones(batch, columns, dtype=torch.bool, device=device) + return torch.rand(batch, width, device=device)[:, :columns] < transform.seg_unit_prob + + +def keep_units(transform: RegionSelecting, units: RegionUnits, batch: int, device: torch.device) -> Tensor: + """`[batch, n_units]` of which units take part, with no reference to the in/out mode. + + For the transforms whose own work is already per region -- `RandomRedistributeSegGPU` + adds a Gaussian per region and then sums them -- rather than a composite over a + selected area. There is no "inside or outside" to choose there, only which of the + regions it already iterates should contribute, so this is the half of + :func:`region_alpha` that applies. + + The fall-back rule applies as it does there: if the draw kept nothing, keep + everything, so a sample never silently loses the whole augmentation. + """ + width = UNIT_DRAW_WIDTH[transform.seg_unit] + transform.bg_em_clusters[1] + columns = min(units.n_units, width) + keep = _keep_draw(transform, batch, columns, width, device) + keep = keep | ~keep.any(dim=1, keepdim=True) + if units.n_units > columns: + # The fold: units past the draw width share the last column's fate, as they do in + # `region_alpha`. + keep = torch.cat([keep, keep[:, -1:].expand(batch, units.n_units - columns)], dim=1) + return keep + + +def region_alpha(transform: RegionSelecting, units: RegionUnits, modes: Tensor) -> Tensor: + """The `[B, 1, *spatial]` composite weight: 1 where the transform applies, 0 where not. + + Built as one gather off the unit id map. Nothing per-unit is materialised: a noisy + mask under `seg_unit="component"` can hold hundreds of components, and + `[2, 300, 128, 128, 128]` of bool is 1.2 GB. + + The keep/drop draw is a fixed-width block (:data:`UNIT_DRAW_WIDTH`) so the number of + uniforms it spends depends on the configuration and not on how many labels this + particular patch happens to contain -- otherwise every transform after this one would + diverge between a four-label patch and a five-label one. Unit ids at or above the + width fold into the last column, which ties those units' fates together rather than + reading past the draw. + + The fall-back rule -- **if a side kept nothing, keep all of it** -- is applied per + side, and costs no extra draw. It is what makes the degenerate cases come out right: + an empty mask has no occupied foreground unit, so `MODE_IN` is a no-op and `MODE_OUT` + covers everything, which is exactly what the binary version did. + """ + device = units.ids.device + width = UNIT_DRAW_WIDTH[transform.seg_unit] + transform.bg_em_clusters[1] + columns = min(units.n_units, width) + column = units.ids.clamp(0, columns - 1) + + keep = _keep_draw(transform, modes.numel(), columns, width, device) + # Background units come first, and only the high end is folded, so a column is + # foreground exactly when its index is at or past `n_background`. + is_foreground = torch.arange(columns, device=device) >= units.n_background + wants_foreground = (modes == MODE_IN) | (modes == MODE_BOTH) + wants_background = (modes == MODE_OUT) | (modes == MODE_BOTH) + wanted = torch.where(is_foreground.unsqueeze(0), wants_foreground.unsqueeze(1), wants_background.unsqueeze(1)) + + selected = keep & wanted + for side in (is_foreground, ~is_foreground): + on_this_side = selected & side + nothing_kept = ~on_this_side.any(dim=1, keepdim=True) + selected = selected | (nothing_kept & wanted & side) + + alpha = torch.gather(selected, 1, column.flatten(1)).view(column.shape) + everywhere = (modes == MODE_ALL).view(-1, *([1] * (alpha.dim() - 1))) + return (alpha | everywhere).to(torch.float32) + + +def feather_alpha(alpha: Tensor, sigma: float) -> Tensor: + """`alpha` with its boundaries turned from a step into a Gaussian ramp. + + Reflect padding, which is `gaussian_blur3d`'s default and the reason to use it rather + than a local copy: zero padding pulls the weight toward 0 at the patch border, so a + unit touching the border would get a soft edge that is an artefact of the patch rather + than of the segmentation. `gpu/fromSeg.py` carries a note about exactly that bug in + the local blur it replaced. + + Every boundary in `alpha` is feathered, not only foreground against background: under + `seg_unit="label"` or `"component"` that includes the boundary between a kept unit and + a dropped neighbour, and under `bg_split="em"` the boundaries between clusters -- which + is what stops an EM split being visible as a hard intensity edge. + + The ramp is centred on the boundary and passes through 0.5 there, so feathering + shrinks the full-strength region slightly. Dilating by `sigma` first would grow the + ramp outward instead and keep each unit at full strength; that is one extra + `max_pool3d` and a second parameter, and worth adding only if the shrink turns out to + matter. + + The `dim` guard is defensive only: `apply_region_selection` rejects anything but three + spatial dimensions before it gets here, because the blur is 3-D. + """ + if sigma <= 0.0 or alpha.dim() != 5: + return alpha + return gaussian_blur3d(alpha, sigma).clamp(0.0, 1.0) + + +# ── the public entry point ─────────────────────────────────────────────────── + + +def _rescaled_into(source: Tensor, target: Tensor, reduce_dims: tuple[int, ...]) -> Tensor: + """`source` linearly remapped onto `target`'s per-sample min-max range.""" + source_min = torch.amin(source, dim=reduce_dims, keepdim=True) + source_max = torch.amax(source, dim=reduce_dims, keepdim=True) + target_min = torch.amin(target, dim=reduce_dims, keepdim=True) + target_max = torch.amax(target, dim=reduce_dims, keepdim=True) + return (source - source_min) / (source_max - source_min + 1e-8) * (target_max - target_min) + target_min + + +def apply_region_selection( + transform: RegionSelecting, + orig: Tensor, + transformed: Tensor, + seg: Tensor | None, + normalize: bool = False, +) -> Tensor: + """Composite `transformed` over `orig` wherever the drawn region says it applies. + + `orig` and `transformed` are either `[N, *spatial]` -- one channel of a batch, which is + what most `apply_transform` methods work in -- or `[*spatial]` for the transforms that + assemble their result one sample at a time. `seg` is the segmentation in any of the + three layouts :func:`label_map_from` reads. + + `normalize` rescales `transformed` into `orig`'s range first. That matters for the + transforms whose output lives in a different space from their input: a SynthSeg or + PALETTE synthesis is in `[0, 1]` while the patch it is composited onto is z-scored, so + a raw composite would put an intensity step on the mask boundary that has nothing to do + with the augmentation. The parameter existed on the function this replaces and had no + caller; these are the callers it was written for. + + Returns `transformed` untouched, having drawn nothing, when there is no segmentation or + when neither probability can fire. + """ + if seg is None or not region_selection_active(transform): + return transformed + if orig.dim() not in (3, 4): + # The contract `_apply_region_mode` had: a channel of a batch, or one sample's + # volume. Three spatial dimensions either way, because the feather blur is 3-D. + raise ValueError(f"Only 4D [N, D, H, W] and 3D [D, H, W] channels are supported. Got {orig.dim()}D.") + + batched = orig.dim() == 4 + reference = orig if batched else orig.unsqueeze(0) + modes = choose_region_modes(transform, reference.shape[0], orig.device) + assert modes is not None # region_selection_active was just checked + sigma = draw_feather(transform) + + units = region_units(seg, reference.unsqueeze(1), unit_spec(transform)) + if units.ids.shape[0] != reference.shape[0]: + # kornia slices the image down to the rows that drew below `p`, and + # `_seg_for_applied_rows` slices the mask to match, so these agree by + # construction. Said out loud because the alternative is a broadcast that + # silently pairs one sample's units with another sample's intensities. + raise ValueError( + f"{type(transform).__name__}: the segmentation has {units.ids.shape[0]} rows but the image has " + f"{reference.shape[0]}; params['seg'] must be sliced to the rows apply_transform was handed." + ) + alpha = feather_alpha(region_alpha(transform, units, modes), sigma) + # `[B, 1, *spatial]` down to the caller's shape. The channel axis exists only so the + # blur above sees a 5-D volume. + weight = alpha.squeeze(1).to(transformed.dtype) + if not batched: + weight = weight.squeeze(0) + + if normalize: + reduce_dims = tuple(range(1, orig.dim())) if batched else tuple(range(orig.dim())) + transformed = _rescaled_into(transformed, orig, reduce_dims) + return weight * transformed + (1.0 - weight) * orig + + +def _select_and_check( + transform: RegionSelecting, + orig: Tensor, + x: Tensor, + seg_mask: Tensor | None, + note: str = "", + normalize: bool = False, +) -> Tensor | None: + """Apply region selection, then reject the result if it went non-finite. + + Returns None when the channel should be left as it was -- the "Final safety" check + that closed all nine of these loops identically. + + Takes the transform rather than its region attributes: every call site read exactly + the same ones off `self`, and spelling them out made the call longer than the code it + replaces. The selection itself lives in `gpu/regions.py`, which is also what decides + whether it runs at all -- a transform with both probabilities at zero gets its input + back having drawn nothing. + """ + x = apply_region_selection(transform, orig, x, seg_mask, normalize=normalize) + # One `isfinite(...).all()`, not `isnan().any() or isinf().any()`: each of those + # reads a reduction back to the host, and the `or` forces the second on the common + # path. Runs for every channel of every region-selecting transform. + if not bool(torch.isfinite(x).all()): + print(f"Warning nan: {type(transform).__name__}{note}", flush=True) + return None + return x diff --git a/smauglab/transforms/synthseg/transforms.py b/smauglab/transforms/synthseg/transforms.py index 735cc2b..ec6f201 100644 --- a/smauglab/transforms/synthseg/transforms.py +++ b/smauglab/transforms/synthseg/transforms.py @@ -21,6 +21,7 @@ import json import os +from collections.abc import Sequence from typing import Any import torch @@ -29,6 +30,16 @@ 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.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, +) from smauglab.transforms.synthseg.generator import SynthSegGenerator # Keys understood from the JSON config / kwargs, forwarded to SynthSegGenerator. @@ -124,6 +135,13 @@ def __init__( self, apply_to_channel: list[int] | None = None, same_on_batch: bool = False, + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 0.5, @@ -134,8 +152,15 @@ def __init__( 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. + # `_GENERATOR_KEYS` on the way to SynthSegGenerator, so a mix or region parameter + # arriving that way would be dropped without a word. + # + # `bg_em_clusters` and the generator's own `em_background_clusters_range` are + # different features one word apart. The generator's splits the *label map it + # synthesises from*, to cope with an incomplete segmentation; this one splits the + # background into units that `out_seg` can land on part of. Setting one does not + # set the other. + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) gen_kwargs = _filter_generator_kwargs(kwargs) # Intensity-only: never deform/flip internally (geometry comes from the @@ -165,10 +190,16 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ 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) + # `_select_and_check` subsumes the hand-written guard this replaces, and costs + # one reduction where `isnan().any() or isinf().any()` read two back to the + # host. `normalize=True` because the synthesis is a [0, 1] GMM render while the + # patch it is composited onto is z-scored; without it a confined synthesis would + # leave an intensity step on the mask boundary, which is a feature the network + # would happily learn. + checked = _select_and_check(self, input[:, c], x, seg, normalize=True) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input diff --git a/unit_tests/helpers.py b/unit_tests/helpers.py index 8c12622..2e7a035 100644 --- a/unit_tests/helpers.py +++ b/unit_tests/helpers.py @@ -49,6 +49,26 @@ def tiny_seg(self) -> torch.Tensor: seg[:, :, 6:18, 6:18, 6:18] = 1.0 return seg + def multi_label_seg(self) -> torch.Tensor: + """Three label values, one of them in two disconnected pieces. + + `tiny_seg` cannot express either new granularity: a single binary cube has one + label value and one connected component, so `seg_unit="label"` and + `seg_unit="component"` are both indistinguishable from `"whole"` on it. Label 1 + here is two separated blocks, which is the only way to tell "component" from + "label" apart, and the three values are what makes "label" more than one unit. + + Blocks rather than noise, so a test can name the voxels it expects to change. + Laid out with background between them, so dilation and feathering have somewhere + to go and the components stay separate under a one-voxel structuring element. + """ + seg = torch.zeros(*VOLUME_SHAPE, dtype=torch.float32) + seg[:, :, 2:6, 2:6, 2:6] = 1.0 # label 1, first component + seg[:, :, 2:6, 2:6, 16:20] = 1.0 # label 1, second component + seg[:, :, 10:14, 10:14, 10:14] = 2.0 + seg[:, :, 16:20, 16:20, 4:8] = 3.0 + return seg + def empty_seg(self) -> torch.Tensor: """A mask with no foreground at all -- every voxel background. diff --git a/unit_tests/test_connected_components.py b/unit_tests/test_connected_components.py new file mode 100644 index 0000000..0be50f5 --- /dev/null +++ b/unit_tests/test_connected_components.py @@ -0,0 +1,225 @@ +"""`connected_components_3d` against scipy, which is the only reference that matters. + +This is the one new algorithm in the region-selection feature, so it is checked against +`scipy.ndimage.label` rather than against its own output. scipy is an unconditional +runtime dependency (see pyproject.toml), so there is no skip guard. + +Two things the tests pin that are easy to get wrong and silent when wrong: + +* **The connectivity is 6, not 26.** The obvious GPU implementation is one + `max_pool3d(kernel=3)`, which is a 3x3x3 neighbourhood and therefore merges blobs that + touch only at a corner. `scipy.ndimage.label`'s default structuring element is + `generate_binary_structure(3, 1)`, the six-neighbour cross, and that is also the + conventional reading of "connected components of a label". +* **Components do not cross label boundaries**, and that holds in a single propagation + pass rather than one pass per label value. A spine patch carries five to twenty-five + labels, so the difference is not academic. + +Partitions are compared, not ids: the numbering is an implementation detail and only the +grouping is the contract. +""" + +from __future__ import annotations + +import numpy as np +import torch +from scipy import ndimage +from torch.nn import functional as F + +from smauglab.transforms.gpu.regions import connected_components_3d +from unit_tests.helpers import SmaugLabTestCase + +#: scipy's default structuring element, spelled out because the default is the whole point. +SIX_CONNECTED = ndimage.generate_binary_structure(3, 1) + + +def as_batch(mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """A `[D, H, W]` bool mask as the `(labels, foreground)` pair the function takes.""" + labels = mask.long().view(1, 1, *mask.shape) + return labels, labels > 0 + + +def blobs(size: int = 20, seed: int = 0) -> torch.Tensor: + """A smoothed random volume thresholded at its mean: many components, odd shapes.""" + torch.manual_seed(seed) + smoothed = F.avg_pool3d(F.pad(torch.rand(1, 1, size, size, size), [3] * 6, mode="reflect"), 7, 1, 0) + return smoothed[0, 0] > smoothed.mean() + + +def same_partition(ours: np.ndarray, theirs: np.ndarray, mask: np.ndarray) -> bool: + """Whether two id maps group the masked voxels identically. + + True when the (ours, theirs) pairs present form a bijection: every one of our + components maps to exactly one of theirs and vice versa. Comparing ids directly would + only test that two numberings happen to agree. + """ + pairs = set(zip(ours[mask].tolist(), theirs[mask].tolist())) + ours_seen = {a for a, _ in pairs} + theirs_seen = {b for _, b in pairs} + return len(pairs) == len(ours_seen) == len(theirs_seen) + + +class TestAgreesWithScipy(SmaugLabTestCase): + def test_random_blobs_partition_identically(self): + for seed in range(6): + with self.subTest(seed=seed): + mask = blobs(seed=seed) + labels, foreground = as_batch(mask) + + ids, count = connected_components_3d(labels, foreground) + reference, expected = ndimage.label(mask.numpy(), structure=SIX_CONNECTED) + + self.assertEqual(count, expected, "a different number of components than scipy found") + self.assertTrue( + same_partition(ids[0, 0].numpy(), reference, mask.numpy()), + "the same count but a different grouping", + ) + + def test_a_multi_label_volume_matches_scipy_run_per_label(self): + """One propagation pass must equal the sum of scipy's per-label runs.""" + for seed in range(4): + with self.subTest(seed=seed): + torch.manual_seed(seed) + smoothed = F.avg_pool3d(F.pad(torch.rand(1, 1, 20, 20, 20), [3] * 6, mode="reflect"), 7, 1, 0)[0, 0] + labels = torch.zeros(20, 20, 20, dtype=torch.long) + labels[smoothed > smoothed.mean()] = 1 + labels[smoothed > smoothed.mean() + 0.6 * smoothed.std()] = 2 + batched = labels.view(1, 1, 20, 20, 20) + + _ids, count = connected_components_3d(batched, batched > 0) + expected = sum(ndimage.label((labels == value).numpy(), structure=SIX_CONNECTED)[1] for value in (1, 2)) + + self.assertEqual(count, expected) + + +class TestTheConnectivityIsSix(SmaugLabTestCase): + def test_two_voxels_touching_at_a_corner_are_two_components(self): + """The test that pins 6- over 26-connectivity. A single max_pool3d merges these.""" + labels = torch.zeros(1, 1, 4, 4, 4, dtype=torch.long) + labels[0, 0, 1, 1, 1] = 1 + labels[0, 0, 2, 2, 2] = 1 + + _ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 2, "corner-touching voxels were merged, so the neighbourhood is 26 not 6") + self.assertEqual(ndimage.label((labels[0, 0] > 0).numpy(), structure=SIX_CONNECTED)[1], 2) + + def test_two_voxels_touching_at_a_face_are_one_component(self): + labels = torch.zeros(1, 1, 4, 4, 4, dtype=torch.long) + labels[0, 0, 1, 1, 1] = 1 + labels[0, 0, 1, 1, 2] = 1 + + _ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + + +class TestLabelBoundariesAreRespected(SmaugLabTestCase): + def test_two_adjacent_slabs_of_different_labels_stay_apart(self): + """They touch across a face, so only the label gate keeps them separate.""" + labels = torch.zeros(1, 1, 8, 8, 8, dtype=torch.long) + labels[0, 0, :, :, :4] = 1 + labels[0, 0, :, :, 4:] = 2 + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 2, "the propagation crossed a label boundary") + self.assertEqual(ids[0, 0, :, :, :4].unique().numel(), 1) + self.assertEqual(ids[0, 0, :, :, 4:].unique().numel(), 1) + self.assertNotEqual(int(ids[0, 0, 0, 0, 0]), int(ids[0, 0, 0, 0, 7])) + + +class TestLongThinStructuresConverge(SmaugLabTestCase): + """The iteration-count test: pointer jumping is what makes these logarithmic. + + A serpentine ribbon has a geodesic length far beyond its Euclidean extent, so plain + neighbour propagation would need one round per voxel along it. + """ + + def test_a_serpentine_ribbon_is_one_component(self): + size = 24 + labels = torch.zeros(1, 1, size, size, size, dtype=torch.long) + row = 0 + for plane in range(0, size, 2): + labels[0, 0, plane, row, :] = 1 + labels[0, 0, plane : min(plane + 2, size), row, size - 1 if plane % 4 == 0 else 0] = 1 + row = min(row + 2, size - 1) + + ids, count = connected_components_3d(labels, labels > 0) + reference, expected = ndimage.label((labels[0, 0] > 0).numpy(), structure=SIX_CONNECTED) + + self.assertEqual(count, expected) + self.assertTrue(same_partition(ids[0, 0].numpy(), reference, (labels[0, 0] > 0).numpy())) + + +class TestDegenerateMasks(SmaugLabTestCase): + def test_an_empty_mask_has_no_components_and_does_not_raise(self): + labels = torch.zeros(1, 1, 5, 5, 5, dtype=torch.long) + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertTrue(bool((ids == -1).all()), "background must be -1, not a unit id") + self.assertGreaterEqual(count, 1, "the count sizes a draw and must never be zero") + + def test_a_full_volume_is_one_component(self): + labels = torch.ones(1, 1, 5, 5, 5, dtype=torch.long) + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + self.assertEqual(ids.unique().tolist(), [0]) + + def test_a_single_voxel_is_one_component(self): + labels = torch.zeros(1, 1, 5, 5, 5, dtype=torch.long) + labels[0, 0, 2, 2, 2] = 1 + + _ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + + +class TestPerSampleNumbering(SmaugLabTestCase): + """Ids are 0-based within each sample, because the keep draw is per sample.""" + + def test_each_sample_is_numbered_from_zero(self): + labels = torch.zeros(2, 1, 6, 6, 6, dtype=torch.long) + labels[0, 0, 0, 0, 0] = 1 + labels[1, 0, 0, 0, 0] = 1 + labels[1, 0, 2, 2, 2] = 1 + labels[1, 0, 4, 4, 4] = 1 + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 3, "the count is the largest per-sample count, which sizes the draw") + self.assertEqual(ids[0, 0][labels[0, 0] > 0].tolist(), [0]) + self.assertEqual(sorted(ids[1, 0][labels[1, 0] > 0].tolist()), [0, 1, 2]) + + def test_two_samples_rooted_at_the_same_index_are_not_merged(self): + """Seeds are globally unique for exactly this reason.""" + labels = torch.zeros(2, 1, 4, 4, 4, dtype=torch.long) + labels[:, 0, 0, 0, 0] = 1 + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + self.assertEqual(int(ids[0, 0, 0, 0, 0]), 0) + self.assertEqual(int(ids[1, 0, 0, 0, 0]), 0) + + +class TestIdsAreUsableAsIndices(SmaugLabTestCase): + """`region_alpha` gathers a per-unit draw with these, so they must be dense.""" + + def test_ids_are_contiguous_from_zero(self): + mask = blobs(seed=3) + labels, foreground = as_batch(mask) + + ids, count = connected_components_3d(labels, foreground) + + present = ids[foreground].unique().tolist() + self.assertEqual(present, list(range(count))) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/unit_tests/test_redistribute_autocast.py b/unit_tests/test_redistribute_autocast.py index cef98bc..e9d7934 100644 --- a/unit_tests/test_redistribute_autocast.py +++ b/unit_tests/test_redistribute_autocast.py @@ -77,7 +77,7 @@ def test_it_moves_the_patch(self): with self.subTest(seed=seed): torch.manual_seed(seed) img, seg = pair() - transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + transform = RandomRedistributeSegGPU(confine_prob=1.0, p=1.0, p_batch=1.0) with torch.autocast("cpu", torch.float16): out = transform.apply_transform(img.clone(), {"seg": seg}, {}) @@ -91,7 +91,7 @@ def test_it_moves_the_patch(self): def test_autocast_does_not_change_the_result(self): """The moments are the only thing autocast could touch here, and they are pinned.""" img, seg = pair() - transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + transform = RandomRedistributeSegGPU(confine_prob=1.0, p=1.0, p_batch=1.0) torch.manual_seed(0) plain = transform.apply_transform(img.clone(), {"seg": seg}, {}) diff --git a/unit_tests/test_redistribute_in_seg.py b/unit_tests/test_redistribute_confine.py similarity index 71% rename from unit_tests/test_redistribute_in_seg.py rename to unit_tests/test_redistribute_confine.py index d3e19c0..0437711 100644 --- a/unit_tests/test_redistribute_in_seg.py +++ b/unit_tests/test_redistribute_confine.py @@ -1,13 +1,19 @@ -"""`in_seg` must confine the redistribution to the segmentation. +"""`confine_prob` must confine the redistribution to the segmentation. + +The parameter was called `in_seg` until this file was renamed with it: everywhere else +in the package `in_seg` is the probability of applying the whole transform inside the +segmentation, while here it decides whether each label's Gaussian is restricted to that +label's own voxels. Two different questions under one name, and twenty-eight shipped +configs set the old one. `aug_redistribute_seg` builds a `to_add` field -- zero outside the labels on the -`in_seg` branch -- and then normalised it with +confined branch -- and then normalised it with img += 2 * (to_add - to_add.min()) / (to_add.max() - to_add.min() + 1e-6) A min-max rescale moves zero. A voxel where nothing was redistributed picked up `-2 * min / range`, which is non-zero whenever `min < 0` -- and it always is, -because the per-label scale is `2 * rand - 1` on [-1, 1]. So `in_seg` restricted +because the per-label scale is `2 * rand - 1` on [-1, 1]. So confining restricted where the redistribution was *computed* but not where it landed: the whole patch got a uniform DC offset. @@ -36,21 +42,21 @@ def normalised(img: torch.Tensor) -> torch.Tensor: return (img - img.min()) / (img.max() - img.min()).clamp_min(1e-6) -class TestInSegStaysInsideTheSegmentation(SmaugLabTestCase): +class TestConfiningStaysInsideTheSegmentation(SmaugLabTestCase): def test_nothing_changes_outside_the_labels(self): for seed in range(5): with self.subTest(seed=seed): torch.manual_seed(seed) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=1.0) delta = out - normalised(img) self.assertAlmostEqual( float(delta[seg == 0].abs().max()), 0.0, places=6, - msg="in_seg applied a DC offset to the background", + msg="confining applied a DC offset to the background", ) def test_something_still_changes_inside_the_labels(self): @@ -58,7 +64,7 @@ def test_something_still_changes_inside_the_labels(self): torch.manual_seed(0) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=1.0) delta = out - normalised(img) self.assertGreater(float(delta[seg == 1].abs().max()), 0.1) @@ -68,17 +74,17 @@ def test_the_peak_amplitude_is_unchanged(self): torch.manual_seed(0) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=1.0) delta = (out - normalised(img)).abs().max() self.assertLessEqual(float(delta), 2.0 + 1e-4) def test_the_whole_image_branch_still_redistributes(self): - """`in_seg=0` takes the other branch, which must keep working.""" + """`confine_prob=0` takes the other branch, which must keep working.""" torch.manual_seed(0) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=0.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=0.0) delta = out - normalised(img) self.assertGreater(float(delta.abs().max()), 0.1) @@ -89,6 +95,6 @@ def test_a_constant_patch_is_still_finite(self): torch.manual_seed(0) _, seg = pair() - out, _ = aug_redistribute_seg(torch.full(SHAPE, -2.709), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(torch.full(SHAPE, -2.709), seg, confine_prob=1.0) self.assertTrue(bool(torch.isfinite(out).all())) diff --git a/unit_tests/test_region_and_stats.py b/unit_tests/test_region_and_stats.py index 896aff1..5b49589 100644 --- a/unit_tests/test_region_and_stats.py +++ b/unit_tests/test_region_and_stats.py @@ -21,9 +21,15 @@ RandomHistogramEqualizationGPU, RandomScharrGPU, RandomSqrtGPU, - _apply_region_mode, _foreground, ) +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + apply_region_selection, +) from unit_tests.helpers import SmaugLabTestCase, first_output @@ -62,7 +68,14 @@ def test_background_stays_background(self): class TestApplyRegionMode(SmaugLabTestCase): - """The behaviour a config actually asks for when it sets in_seg or out_seg.""" + """The behaviour a config actually asks for when it sets in_seg or out_seg. + + The mode is no longer passed in as a string: `in_seg` and `out_seg` are drawn per + sample, so 1.0/0.0 is how a test pins mode "in". `seg_unit_prob=1.0` keeps every unit, + which is what makes these assertions about *where* rather than about which units the + keep draw happened to pick, and `seg_feather=0` keeps the boundary a step so the + changed-voxel set is exactly the selected region. + """ def setUp(self): super().setUp() @@ -74,29 +87,43 @@ def setUp(self): self.mask[0, 0, :2, :2, :2] = 1.0 self.inside = (slice(None), slice(0, 2), slice(0, 2), slice(0, 2)) + def _select(self, **kwargs): + return apply_region_selection(_Stub(**kwargs), self.orig, self.transformed, self.mask) + def test_mode_in_changes_only_the_masked_voxels(self): - out = _apply_region_mode(self.orig, self.transformed, self.mask, "in") + out = self._select(in_seg=1.0) self.assertTrue(bool((out[self.inside] == 1.0).all()), "'in' did not apply the transform inside the mask") self.assertEqual(int(out.sum()), 8, "'in' leaked outside the mask") def test_mode_out_changes_only_the_unmasked_voxels(self): - out = _apply_region_mode(self.orig, self.transformed, self.mask, "out") + out = self._select(out_seg=1.0) self.assertTrue(bool((out[self.inside] == 0.0).all()), "'out' applied the transform inside the mask") self.assertEqual(int(out.sum()), 64 - 8) def test_in_and_out_partition_the_volume(self): - inside = _apply_region_mode(self.orig, self.transformed, self.mask, "in") - outside = _apply_region_mode(self.orig, self.transformed, self.mask, "out") + inside = self._select(in_seg=1.0) + outside = self._select(out_seg=1.0) self.assertTrue(torch.equal(inside + outside, self.transformed)) + def test_both_fired_is_a_mixture_and_not_the_whole_patch(self): + """The one semantic change: `in_seg` and `out_seg` both firing used to mean "all". + + At `seg_unit_prob=1.0` every unit of both sides is kept, so "both" does cover the + whole patch -- the difference from "all" only shows up below that, where the two + sides are sampled independently. + """ + covered = apply_region_selection(_Stub(in_seg=1.0, out_seg=1.0, seg_unit_prob=1.0), self.orig, self.transformed, self.mask) + self.assertTrue(torch.equal(covered, self.transformed)) + def test_mode_all_ignores_the_mask(self): - out = _apply_region_mode(self.orig, self.transformed, self.mask, "all") + """Neither probability firing. The mask is read but nothing is confined.""" + out = self._select(in_seg=0.0, out_seg=0.0) self.assertTrue(torch.equal(out, self.transformed)) def test_a_missing_mask_means_apply_everywhere(self): - out = _apply_region_mode(self.orig, self.transformed, None, "in") + out = apply_region_selection(_Stub(in_seg=1.0), self.orig, self.transformed, None) self.assertTrue(torch.equal(out, self.transformed)) def test_the_unbatched_3d_path_behaves_the_same(self): @@ -105,12 +132,28 @@ def test_the_unbatched_3d_path_behaves_the_same(self): mask = torch.zeros(2, 4, 4, 4) mask[0, :2, :2, :2] = 1.0 - out = _apply_region_mode(orig, transformed, mask, "in") + out = apply_region_selection(_Stub(in_seg=1.0), orig, transformed, mask) self.assertEqual(int(out.sum()), 8) def test_an_unsupported_rank_is_rejected(self): with self.assertRaises(ValueError): - _apply_region_mode(torch.rand(2, 2), torch.rand(2, 2), torch.rand(1, 2, 2), "in") + apply_region_selection(_Stub(in_seg=1.0), torch.rand(2, 2), torch.rand(2, 2), torch.rand(1, 2, 2)) + + def test_a_feathered_boundary_is_a_ramp_not_a_step(self): + """`seg_feather` is the whole point of the alpha being a float and not a bool.""" + step = self._select(in_seg=1.0) + ramp = apply_region_selection(_Stub(in_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.mask) + + self.assertEqual(sorted(set(step.flatten().tolist())), [0.0, 1.0], "the unfeathered composite is a step") + between = (ramp > 1e-6) & (ramp < 1.0 - 1e-6) + self.assertGreater(int(between.sum()), 0, "a feathered boundary must put voxels strictly between the two") + self.assertTrue(bool((ramp >= 0.0).all() and (ramp <= 1.0).all()), "the alpha left [0, 1]") + + def test_feathered_in_and_out_still_sum_to_the_transform(self): + """blur(m) + blur(1 - m) == blur(1) == 1, for a normalised kernel and reflect padding.""" + inside = apply_region_selection(_Stub(in_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.mask) + outside = apply_region_selection(_Stub(out_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.mask) + self.assertTrue(torch.allclose(inside + outside, self.transformed, atol=1e-5)) class TestRegionModeReachesTheRealTransforms(SmaugLabTestCase): @@ -237,7 +280,7 @@ def test_a_2d_image_does_not_raise(self): seg = torch.zeros(12, 12) seg[3:8, 3:8] = 1.0 - out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(image.clone(), seg, confine_prob=1.0) self.assertEqual(tuple(out.shape), tuple(image.shape)) self.assertTrue(bool(np.isfinite(out.numpy()).all())) @@ -249,18 +292,31 @@ def test_the_3d_path_is_unchanged(self): seg = torch.zeros(10, 10, 10) seg[2:6, 2:6, 2:6] = 1.0 - out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(image.clone(), seg, confine_prob=1.0) self.assertEqual(tuple(out.shape), tuple(image.shape)) class _Stub: - """Stands in for a transform: _select_and_check only reads these three attributes.""" + """Stands in for a transform: the region helpers only read these attributes. + + Built from `gpu/regions.py`'s own "off" defaults rather than from literals, so this + also fails if one of those constants disappears. + """ - def __init__(self, in_seg: float = 0.0, out_seg: float = 0.0, mix_in_out: bool = False): - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + def __init__(self, **overrides): + defaults = { + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": DEFAULT_SEG_UNIT, + "seg_unit_prob": 1.0, + "bg_split": DEFAULT_BG_SPLIT, + "bg_em_clusters": DEFAULT_BG_EM_CLUSTERS, + "seg_feather": (DEFAULT_SEG_FEATHER, DEFAULT_SEG_FEATHER), + } + defaults.update(overrides) + for name, value in defaults.items(): + setattr(self, name, value) class TestExtractedTailHelpers(SmaugLabTestCase): diff --git a/unit_tests/test_region_selection.py b/unit_tests/test_region_selection.py new file mode 100644 index 0000000..a66f24d --- /dev/null +++ b/unit_tests/test_region_selection.py @@ -0,0 +1,682 @@ +"""Where a transform applies, at what granularity, and how sharply it stops. + +`in_seg` and `out_seg` predate this; what they did not express is most of what they +looked like they expressed. The parameters and the four defects they replace are set out +in `smauglab/transforms/gpu/regions.py`; this file is what holds the behaviour. + +The registry-driven cases matter most. The surface has to reach every transform that +reads a segmentation, and a transform added later must not be able to arrive without it +-- or, worse, with the parameters wired into its constructor and never read. + +Three invariants here are about the **RNG stream** rather than about pixels, and all +three are properties the code this replaces broke: + +* a knob that is off spends no random numbers, so a sweep is comparable with its own + baseline at a fixed seed; +* the number of uniforms a draw spends depends on the configuration and never on the + data, so a four-label patch and a five-label one leave the stream in the same place; +* the EM background split is forked, so a cache hit and a cache miss agree. +""" + +from __future__ import annotations + +import gc + +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 import regions +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.contrast import RandomScharrGPU +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + UnitSpec, + apply_region_selection, + configure_regions, + region_units, +) +from unit_tests.helpers import SmaugLabTestCase, domain_bank_missing, first_output + +BATCH = 2 +SHAPE = (BATCH, 1, 24, 24, 24) + +#: `RandomRedistributeSegGPU` is the one transform that deliberately does not take +#: `in_seg` / `out_seg`: it used to accept `in_seg` with a different meaning, and had it +#: kept the name while changing it, twenty-eight shipped configs would have flipped +#: silently. See its class docstring. +NO_MODE_PARAMETERS = {"RandomRedistributeSegGPU"} + +#: The five that are not about the in/out mode and therefore reach every segmentation-aware +#: transform, including the one above. +GRANULARITY_PARAMETERS = ("seg_unit", "seg_unit_prob", "bg_split", "bg_em_clusters", "seg_feather") + + +def region_aware(): + """Every registered GPU transform that takes the region surface.""" + return [entry for entry in registry.entries(Backend.GPU) if "seg_unit" in registry.accepted_params(entry)] + + +def transfer_augmentations(): + return registry.entries(Backend.GPU, group=AugType.TA) + + +class _Stub: + """Stands in for a transform. The helpers only read these attributes.""" + + def __init__(self, **overrides): + values = { + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": DEFAULT_SEG_UNIT, + "seg_unit_prob": 1.0, + "bg_split": DEFAULT_BG_SPLIT, + "bg_em_clusters": DEFAULT_BG_EM_CLUSTERS, + "seg_feather": (DEFAULT_SEG_FEATHER, DEFAULT_SEG_FEATHER), + } + values.update(overrides) + for name, value in values.items(): + setattr(self, name, value) + + +class TestTheSurfaceReachesEveryTransform(SmaugLabTestCase): + """Registry-driven, so a transform added later is covered when it is registered.""" + + def test_every_transfer_augmentation_takes_the_granularity_parameters(self): + for entry in transfer_augmentations(): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + for name in GRANULARITY_PARAMETERS: + self.assertIn(name, accepted, f"{entry.name} cannot be given {name}") + + def test_every_transfer_augmentation_but_one_takes_the_mode_parameters(self): + for entry in transfer_augmentations(): + accepted = registry.accepted_params(entry) + expected = entry.name not in NO_MODE_PARAMETERS + with self.subTest(entry=entry.name): + self.assertEqual("in_seg" in accepted, expected) + self.assertEqual("out_seg" in accepted, expected) + + def test_mix_in_out_is_gone_everywhere(self): + for entry in registry.entries(Backend.GPU): + with self.subTest(entry=entry.name): + self.assertNotIn("mix_in_out", registry.accepted_params(entry)) + + def test_the_defaults_mean_off(self): + """A config that sets none of these must get the behaviour it had before.""" + for entry in region_aware(): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + self.assertEqual(accepted["bg_split"].default, DEFAULT_BG_SPLIT) + self.assertEqual(accepted["seg_feather"].default, DEFAULT_SEG_FEATHER) + self.assertEqual(tuple(accepted["bg_em_clusters"].default), DEFAULT_BG_EM_CLUSTERS) + if entry.name in NO_MODE_PARAMETERS: + # Its own defaults reproduce the region set it has always used. + self.assertEqual(accepted["seg_unit"].default, "label") + self.assertEqual(accepted["seg_unit_prob"].default, 1.0) + else: + self.assertEqual(accepted["in_seg"].default, 0.0) + self.assertEqual(accepted["out_seg"].default, 0.0) + self.assertEqual(accepted["seg_unit"].default, DEFAULT_SEG_UNIT) + self.assertEqual(accepted["seg_unit_prob"].default, DEFAULT_SEG_UNIT_PROB) + + def test_the_knobs_are_read_and_not_merely_declared(self): + """The guard against a constructor that accepts a parameter and ignores it. + + Several seeds, because a single draw can legitimately keep every label and then + `seg_unit="label"` *is* `"whole"`. A transform that leaves the volume alone on + this input is skipped rather than failed: there is nothing to confine, which is + a statement about the input and not about the knob. + """ + image = torch.rand(*SHAPE) + seg = self.multi_label_seg().expand(BATCH, -1, -1, -1, -1).contiguous() + for entry in region_aware(): + with self.subTest(entry=entry.name): + skip = domain_bank_missing(entry) if entry.external_asset else None + if skip: + self.skipTest(skip) + extra = dict(entry.smoke_kwargs) + mode = {} if entry.name in NO_MODE_PARAMETERS else {"in_seg": 1.0} + + def run(seed, entry=entry, extra=extra, **kwargs): + torch.manual_seed(seed) + transform = entry.cls(p=1.0, **extra, **kwargs) + torch.manual_seed(seed) + return transform.apply_transform(image.clone(), {"seg": seg.clone()}, transform.flags) + + untouched = all(bool(torch.equal(run(seed, **mode), image)) for seed in range(4)) + if untouched: + self.skipTest(f"{entry.name} leaves this input unchanged, so there is nothing to confine") + + differed = any( + not bool( + torch.equal( + run(seed, **mode, seg_unit="whole", seg_unit_prob=0.5), + run(seed, **mode, seg_unit="label", seg_unit_prob=0.5), + ) + ) + for seed in range(6) + ) + self.assertTrue( + differed, + f"{entry.name} accepts seg_unit but its output does not depend on it", + ) + + +class TestWhereEachModeLands(SmaugLabTestCase): + """`seg_unit_prob=1.0` keeps every unit, so these are assertions about *where*.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + self.seg = self.multi_label_seg() + self.foreground = self.seg[:, 0] > 0 + + def changed(self, **kwargs): + out = apply_region_selection(_Stub(**kwargs), self.orig, self.transformed, self.seg.clone()) + return out != self.orig + + def test_in_covers_exactly_the_labels(self): + self.assertTrue(bool(torch.equal(self.changed(in_seg=1.0), self.foreground))) + + def test_out_covers_exactly_the_background(self): + self.assertTrue(bool(torch.equal(self.changed(out_seg=1.0), ~self.foreground))) + + def test_both_covers_the_patch_when_every_unit_is_kept(self): + self.assertTrue(bool(self.changed(in_seg=1.0, out_seg=1.0).all())) + + def test_neither_is_the_whole_patch_untouched_by_the_mask(self): + out = apply_region_selection(_Stub(), self.orig, self.transformed, self.seg) + self.assertTrue(bool(torch.equal(out, self.transformed))) + + def test_both_is_a_mixture_and_not_merely_all(self): + """At a probability below 1 the two sides are sampled independently. + + This is the one semantic change to `in_seg` / `out_seg`: both firing used to fall + through to "the whole patch", so `in_seg=1, out_seg=1` was indistinguishable from + `in_seg=0, out_seg=0`. + """ + torch.manual_seed(4) + partials = 0 + for _ in range(12): + changed = self.changed(in_seg=1.0, out_seg=1.0, seg_unit="label", seg_unit_prob=0.5) + touches_both = bool((changed & self.foreground).any()) and bool((changed & ~self.foreground).any()) + if touches_both and not bool(changed.all()): + partials += 1 + self.assertGreater(partials, 0, "'both' never produced a mixture, only the whole patch") + + +class TestGranularityKeepsWholeUnits(SmaugLabTestCase): + """A unit is applied or it is not; it is never partially applied.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + self.seg = self.multi_label_seg() + + def assert_whole_units(self, seg_unit): + spec = UnitSpec(seg_unit, DEFAULT_BG_SPLIT, DEFAULT_BG_EM_CLUSTERS) + for seed in range(8): + torch.manual_seed(seed) + seg = self.seg.clone() + units = region_units(seg, self.orig.unsqueeze(1), spec) + out = apply_region_selection(_Stub(in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5), self.orig, self.transformed, seg) + changed = (out != self.orig)[0] + for unit in range(units.n_background, units.n_units): + cells = units.ids[0, 0] == unit + if not bool(cells.any()): + continue + values = changed[cells] + with self.subTest(seg_unit=seg_unit, seed=seed, unit=unit): + self.assertTrue( + bool(values.all()) or bool((~values).all()), + f"{seg_unit} unit {unit} was partially applied", + ) + + def test_label_units_are_whole(self): + self.assert_whole_units("label") + + def test_component_units_are_whole(self): + self.assert_whole_units("component") + + def test_components_split_a_label_that_labels_does_not(self): + """`multi_label_seg`'s label 1 is two separated blocks.""" + by_label = region_units(self.seg.clone(), None, UnitSpec("label", "none", DEFAULT_BG_EM_CLUSTERS)) + by_component = region_units(self.seg.clone(), None, UnitSpec("component", "none", DEFAULT_BG_EM_CLUSTERS)) + self.assertEqual(by_label.n_units, 4, "background plus three label values") + self.assertEqual(by_component.n_units, 5, "the split label contributes one more") + + def test_keeping_every_unit_reduces_to_whole(self): + """`seg_unit_prob=1.0` is how granularity is turned off, and it must be exact.""" + for seg_unit in ("label", "component"): + with self.subTest(seg_unit=seg_unit): + whole = apply_region_selection( + _Stub(in_seg=1.0, seg_unit="whole", seg_unit_prob=1.0), self.orig, self.transformed, self.seg.clone() + ) + granular = apply_region_selection( + _Stub(in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=1.0), self.orig, self.transformed, self.seg.clone() + ) + self.assertTrue(bool(torch.equal(whole, granular))) + + def test_a_single_label_mask_makes_label_the_same_as_whole(self): + """Documented rather than discovered: there is nothing to mix on a binary mask.""" + binary = self.tiny_seg() + whole = apply_region_selection(_Stub(in_seg=1.0, seg_unit="whole"), self.orig, self.transformed, binary.clone()) + by_label = apply_region_selection(_Stub(in_seg=1.0, seg_unit="label"), self.orig, self.transformed, binary.clone()) + self.assertTrue(bool(torch.equal(whole, by_label))) + + +class TestTheBackgroundSplit(SmaugLabTestCase): + """`bg_split="em"` is what makes the background something other than all-or-nothing.""" + + def setUp(self): + super().setUp() + self.seg = self.tiny_seg() + self.foreground = self.seg[:, 0] > 0 + # Bimodal background, so the EM has two modes to find. + torch.manual_seed(1) + self.image = torch.zeros(1, 24, 24, 24) + self.image[0, :12] = 0.2 + self.image[0, 12:] = 0.9 + self.image += torch.rand(1, 24, 24, 24) * 0.02 + self.transformed = torch.full((1, 24, 24, 24), 5.0) + + def changed_background(self, **kwargs): + out = apply_region_selection(_Stub(**kwargs), self.image, self.transformed, self.seg.clone()) + return int(((out != self.image) & ~self.foreground).sum()) + + def test_without_the_split_the_background_is_all_or_nothing(self): + torch.manual_seed(5) + total = int((~self.foreground).sum()) + seen = {self.changed_background(out_seg=1.0, seg_unit_prob=0.5) for _ in range(8)} + self.assertTrue(seen <= {0, total}, f"a partial background without bg_split: {sorted(seen)}") + + def test_with_the_split_part_of_the_background_can_be_used(self): + torch.manual_seed(5) + total = int((~self.foreground).sum()) + counts = [self.changed_background(out_seg=1.0, bg_split="em", seg_unit_prob=0.5) for _ in range(8)] + self.assertTrue( + any(0 < count < total for count in counts), + f"bg_split='em' never produced a partial background: {counts}", + ) + + def test_the_foreground_is_still_untouched_in_mode_out(self): + torch.manual_seed(5) + out = apply_region_selection(_Stub(out_seg=1.0, bg_split="em", seg_unit_prob=0.5), self.image, self.transformed, self.seg) + self.assertTrue(bool(torch.equal(out[self.foreground], self.image[self.foreground]))) + + def test_a_patch_with_no_background_is_not_split(self): + """`em_subdivide_labels` falls back to the largest label when 0 is absent.""" + units = region_units(torch.ones(1, 1, 8, 8, 8), torch.rand(1, 1, 8, 8, 8), UnitSpec("whole", "em", (3, 10))) + self.assertEqual(units.n_background, 1, "the EM fallback split a foreground label") + + def test_a_constant_volume_gives_one_background_unit(self): + """The alternative is clusters whose boundaries are decided by float noise.""" + seg = torch.zeros(1, 1, 10, 10, 10) + seg[0, 0, 4:6, 4:6, 4:6] = 1 + units = region_units(seg, self.constant_volume()[:, :, :10, :10, :10], UnitSpec("whole", "em", (3, 10))) + self.assertEqual(units.n_background, 1) + self.assertTrue(bool(torch.isfinite(units.ids.float()).all())) + + +class TestTheBoundary(SmaugLabTestCase): + """`seg_feather` turns the composite weight from a step into a Gaussian ramp.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + self.seg = self.tiny_seg() + + def composite(self, sigma): + return apply_region_selection(_Stub(in_seg=1.0, seg_feather=(sigma, sigma)), self.orig, self.transformed, self.seg.clone()) + + def test_zero_is_a_hard_step(self): + self.assertEqual(sorted(set(self.composite(0.0).flatten().tolist())), [0.0, 1.0]) + + def test_a_positive_sigma_puts_voxels_between_the_two(self): + ramp = self.composite(1.0) + between = (ramp > 1e-6) & (ramp < 1.0 - 1e-6) + self.assertGreater(int(between.sum()), 0) + self.assertTrue(bool((ramp >= 0.0).all() and (ramp <= 1.0).all()), "the weight left [0, 1]") + + def test_the_ramp_is_monotone_across_the_boundary(self): + ramp = self.composite(1.5)[0, 12, 12, :] + rising = ramp[:12] + self.assertTrue(bool((rising[1:] - rising[:-1] >= -1e-6).all()), rising.tolist()) + + def test_the_interior_of_a_unit_stays_at_full_strength(self): + """The ramp is centred on the boundary, so it does not reach a large unit's middle.""" + self.assertGreater(float(self.composite(1.0)[0, 11:13, 11:13, 11:13].min()), 0.999) + + def test_in_and_out_still_sum_to_the_transform(self): + """blur(m) + blur(1 - m) == blur(1) == 1 for a normalised kernel and reflect padding.""" + inside = self.composite(1.0) + outside = apply_region_selection(_Stub(out_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.seg.clone()) + self.assertTrue(torch.allclose(inside + outside, self.transformed, atol=1e-5)) + + def test_mode_all_is_unaffected(self): + """The blur of a constant weight is that constant.""" + out = apply_region_selection(_Stub(in_seg=1e-12, seg_feather=(2.0, 2.0)), self.orig, self.transformed, self.seg) + self.assertTrue(torch.allclose(out, self.transformed, atol=1e-6)) + + def test_an_empty_mask_is_unaffected(self): + out = apply_region_selection(_Stub(in_seg=1.0, seg_feather=(2.0, 2.0)), self.orig, self.transformed, self.empty_seg()) + self.assertTrue(torch.allclose(out, self.orig, atol=1e-6)) + + +class TestDegenerateMasks(SmaugLabTestCase): + """An empty mask is ~18% of real nnU-Net samples; see `helpers.empty_seg`.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + + def test_in_is_a_no_op_and_out_is_everywhere(self): + for seg_unit in ("whole", "label", "component"): + with self.subTest(seg_unit=seg_unit): + inside = apply_region_selection( + _Stub(in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5), self.orig, self.transformed, self.empty_seg() + ) + outside = apply_region_selection( + _Stub(out_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5), self.orig, self.transformed, self.empty_seg() + ) + self.assertTrue(bool(torch.equal(inside, self.orig)), "'in' applied something with nothing labelled") + self.assertTrue(bool(torch.equal(outside, self.transformed)), "'out' did not cover the whole patch") + + def test_the_unit_count_is_never_zero(self): + """It sizes a draw and a gather; zero is a device-side assert, not a failed batch.""" + for seg_unit in ("whole", "label", "component"): + with self.subTest(seg_unit=seg_unit): + units = region_units(self.empty_seg(), None, UnitSpec(seg_unit, "none", DEFAULT_BG_EM_CLUSTERS)) + self.assertGreaterEqual(units.n_units, 1) + self.assertGreaterEqual(units.n_background, 1) + + +class TestTheDecompositionIsSharedAcrossAPass(SmaugLabTestCase): + """Without this the feature is not trainable: a pipeline runs twenty-odd of these.""" + + def setUp(self): + super().setUp() + self.calls = 0 + self.real = regions._build_units + + def counting(*args, **kwargs): + self.calls += 1 + return self.real(*args, **kwargs) + + regions._build_units = counting + self.addCleanup(setattr, regions, "_build_units", self.real) + + def pipeline(self, *transforms): + return AugmentationSequentialCustom(*transforms, data_keys=["input", "mask"], same_on_batch=True) + + def test_two_transforms_in_one_pass_share_one_decomposition(self): + image = torch.rand(1, 1, 24, 24, 24) + seg = self.multi_label_seg() + pipeline = self.pipeline( + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="component"), + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="component"), + ) + + first_output(pipeline(image.clone(), seg.clone())) + + self.assertEqual(self.calls, 1, "the decomposition was rebuilt per transform") + + def test_a_different_spec_is_a_different_entry(self): + image = torch.rand(1, 1, 24, 24, 24) + seg = self.multi_label_seg() + pipeline = self.pipeline( + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="component"), + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="label"), + ) + + first_output(pipeline(image.clone(), seg.clone())) + + self.assertEqual(self.calls, 2) + + def test_a_new_mask_tensor_is_a_new_decomposition(self): + """How invalidation works: a geometric transform returns a different object.""" + image = torch.rand(1, 1, 24, 24, 24) + spec = UnitSpec("component", "none", DEFAULT_BG_EM_CLUSTERS) + region_units(self.multi_label_seg(), image, spec) + region_units(self.multi_label_seg(), image, spec) + self.assertEqual(self.calls, 2) + + def test_the_cache_dies_with_the_mask(self): + seg = self.multi_label_seg() + region_units(seg, torch.rand(1, 1, 24, 24, 24), UnitSpec("label", "none", DEFAULT_BG_EM_CLUSTERS)) + self.assertTrue(hasattr(seg, "_smauglab_region_units")) + del seg + gc.collect() # nothing else holds it, so there is nothing left to leak + + +class TestTheRngStream(SmaugLabTestCase): + """Three properties without which a seeded sweep is not comparable to its baseline.""" + + def scharr_then_a_draw(self, seed=11, **kwargs): + image = torch.rand(*SHAPE) + seg = self.multi_label_seg().expand(BATCH, -1, -1, -1, -1).contiguous() + transform = RandomScharrGPU(p=1.0, **kwargs) + torch.manual_seed(seed) + transform.apply_transform(image.clone(), {"seg": seg}, transform.flags) + return torch.rand(4) + + def test_region_selection_off_spends_nothing(self): + self.assertTrue( + bool(torch.equal(self.scharr_then_a_draw(), self.scharr_then_a_draw(in_seg=0.0, out_seg=0.0, seg_unit="component"))), + "the region parameters moved the stream with both probabilities at zero", + ) + + def test_the_draw_width_does_not_depend_on_the_data(self): + """Two labels or five, the stream has to land in the same place.""" + two = torch.zeros(*SHAPE) + two[:, :, 2:6, 2:6, 2:6] = 1 + two[:, :, 10:14, 10:14, 10:14] = 2 + five = two.clone() + for offset, value in enumerate((3, 4, 5)): + five[:, :, 20, 20, offset] = value + + def stream(seg, seg_unit): + transform = RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5) + torch.manual_seed(11) + transform.apply_transform(torch.rand(*SHAPE), {"seg": seg.clone()}, transform.flags) + return torch.rand(4) + + for seg_unit in ("label", "component"): + with self.subTest(seg_unit=seg_unit): + self.assertTrue( + bool(torch.equal(stream(two, seg_unit), stream(five, seg_unit))), + f"{seg_unit} spent a number of uniforms that depended on the segmentation", + ) + + def test_keeping_every_unit_spends_nothing(self): + """`seg_unit_prob=1` is how granularity is turned off, so it has to be free.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(in_seg=1.0, seg_unit="whole", seg_unit_prob=1.0), + self.scharr_then_a_draw(in_seg=1.0, seg_unit="component", seg_unit_prob=1.0), + ) + ) + ) + + def test_the_em_split_is_forked_out_of_the_global_stream(self): + """It spends ~100k uniforms on a cache miss and none on a hit.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(out_seg=1.0, bg_split="none"), + self.scharr_then_a_draw(out_seg=1.0, bg_split="em"), + ) + ), + "the EM draw leaked into the global stream, so cache state changes the augmentation", + ) + + def test_the_feather_range_does_not_shift_the_stream_when_it_is_off(self): + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(in_seg=1.0, seg_feather=0.0), + self.scharr_then_a_draw(in_seg=1.0), + ) + ) + ) + + def test_a_degenerate_feather_range_still_spends_its_draw(self): + """Otherwise the stream becomes a function of the range and an A/B is not valid.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(in_seg=1.0, seg_feather=[1.0, 1.0]), + self.scharr_then_a_draw(in_seg=1.0, seg_feather=[0.5, 2.0]), + ) + ) + ) + + +class TestTheModeIsPerSample(SmaugLabTestCase): + """It was one decision for the whole batch at ten of the eleven call sites. + + So `in_seg=0.5` read as "half the batches" where `p` and `mix_prob` both meant "half + the samples". + """ + + def test_all_four_modes_appear_over_a_batch(self): + torch.manual_seed(0) + seen = set() + for _ in range(16): + modes = regions.choose_region_modes(_Stub(in_seg=0.5, out_seg=0.5), BATCH, torch.device("cpu")) + seen.update(modes.tolist()) + self.assertEqual(seen, {regions.MODE_ALL, regions.MODE_IN, regions.MODE_OUT, regions.MODE_BOTH}) + + def test_two_samples_of_one_batch_can_differ(self): + torch.manual_seed(0) + differed = any( + len(set(regions.choose_region_modes(_Stub(in_seg=0.5, out_seg=0.5), BATCH, torch.device("cpu")).tolist())) > 1 + for _ in range(16) + ) + self.assertTrue(differed, "every sample of the batch got the same mode") + + def test_nothing_is_drawn_when_both_probabilities_are_zero(self): + self.assertIsNone(regions.choose_region_modes(_Stub(), BATCH, torch.device("cpu"))) + + +class TestItSurvivesTheConfigPath(SmaugLabTestCase): + """The constructors are the schema, and JSON hands over lists rather than tuples.""" + + def test_a_section_validates_and_reaches_the_transform(self): + section = { + "RandomScharrGPU": { + "in_seg": 0.5, + "out_seg": 0.25, + "seg_unit": "component", + "seg_unit_prob": 0.3, + "bg_split": "em", + "bg_em_clusters": [2, 5], + "seg_feather": [0.5, 1.5], + } + } + self.assertEqual(validate_section(section, Backend.GPU), []) + + built = {type(t).__name__: t for t in build_gpu_pipeline(section)} + scharr = built["RandomScharrGPU"] + self.assertEqual(scharr.seg_unit, "component") + self.assertEqual(scharr.bg_split, "em") + self.assertEqual(scharr.bg_em_clusters, (2, 5)) + self.assertEqual(scharr.seg_feather, (0.5, 1.5)) + + def test_a_bad_value_in_a_section_fails_the_build(self): + """Not at validation -- that only checks keys -- but before the first batch.""" + section = {"RandomScharrGPU": {"seg_unit": "labl"}} + self.assertEqual(validate_section(section, Backend.GPU), []) + with self.assertRaises(ValueError): + build_gpu_pipeline(section) + + def test_the_renamed_parameters_are_rejected_with_guidance(self): + for section, needle in ( + ({"RandomScharrGPU": {"mix_in_out": True}}, "seg_unit"), + ({"RandomRedistributeSegGPU": {"in_seg": 0.25}}, "confine_prob"), + ): + with self.subTest(section=section): + problems = validate_section(section, Backend.GPU) + self.assertEqual(len(problems), 1) + self.assertIn(needle, problems[0]) + + def test_no_shipped_config_still_sets_the_old_names(self): + import json + + from unit_tests.helpers import configs_dir + + for path in sorted(configs_dir().rglob("*.json")): + payload = json.loads(path.read_text()) + for section in payload.values(): + if not isinstance(section, dict): + continue + for name, params in section.items(): + if not isinstance(params, dict): + continue + with self.subTest(config=path.name, augmentation=name): + self.assertNotIn("mix_in_out", params) + if name in {"RandomRedistributeSegGPU", "RedistributeTransform"}: + self.assertNotIn("in_seg", params) + + +class TestTheValuesAreValidated(SmaugLabTestCase): + """At construction, which is where every value check in this package happens.""" + + def configure(self, **kwargs): + class Owner: + pass + + arguments = {"in_seg": 0.0, "out_seg": 0.0} + arguments.update(kwargs) + configure_regions(Owner(), **arguments) + + def test_a_bad_value_names_the_parameter_and_the_class(self): + for kwargs, needle in ( + ({"in_seg": 1.5}, "in_seg"), + ({"out_seg": -0.1}, "out_seg"), + ({"seg_unit": "labl"}, "seg_unit"), + ({"seg_unit_prob": True}, "seg_unit_prob"), + ({"bg_split": "EM"}, "bg_split"), + ({"bg_em_clusters": 5}, "bg_em_clusters"), + ({"bg_em_clusters": [5, 2]}, "bg_em_clusters"), + ({"seg_feather": -1.0}, "seg_feather"), + ): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError) as caught: + self.configure(**kwargs) + self.assertIn(needle, str(caught.exception)) + self.assertIn("Owner", str(caught.exception)) + + def test_the_legal_seg_unit_values_are_listed(self): + with self.assertRaises(ValueError) as caught: + self.configure(seg_unit="nope") + for value in regions.SEG_UNITS: + self.assertIn(value, str(caught.exception)) + + def test_a_feather_above_one_is_accepted(self): + """It is a sigma in voxels, not a weight -- the point of `bounds=`.""" + self.configure(seg_feather=[0.5, 4.0]) + + def test_a_mix_alpha_above_one_is_still_rejected(self): + """The same check, with its own bounds. Widening one must not widen the other.""" + from smauglab.transforms.gpu.mixing import normalise_blend_range + + with self.assertRaises(ValueError): + normalise_blend_range([0.5, 4.0]) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/unit_tests/test_seg_layout.py b/unit_tests/test_seg_layout.py index c3f5e42..2a111a1 100644 --- a/unit_tests/test_seg_layout.py +++ b/unit_tests/test_seg_layout.py @@ -87,7 +87,7 @@ def test_label_map_and_one_hot_give_the_same_image(self): image = torch.rand(*SHAPE) def build(): - return RandomRedistributeSegGPU(p=1.0, in_seg=1.0) + return RandomRedistributeSegGPU(p=1.0, confine_prob=1.0) torch.testing.assert_close(_apply(build, labels, image), _apply(build, onehot, image)) @@ -99,7 +99,7 @@ def test_the_labels_actually_drive_the_result(self): image = torch.rand(*SHAPE) def build(): - return RandomRedistributeSegGPU(p=1.0, in_seg=1.0) + return RandomRedistributeSegGPU(p=1.0, confine_prob=1.0) per_label = _apply(build, labels, image) one_blob = _apply(build, merged, image)