From a23544813d5128229628e7098317e780e2ce6864 Mon Sep 17 00:00:00 2001 From: iback Date: Fri, 21 Aug 2026 06:13:38 +0000 Subject: [PATCH] feat: register every augmentation Each of the 51 augmentation classes now carries an @register(...) decorator naming its backend, its GEO/GE/TA group and the backend-neutral concept it implements. smauglab/transforms/__init__.py imports them all, so importing that package is what populates the registry -- that is what registry.load_all() does. The batchgeneratorsv2 transforms the CPU pipeline composes are third-party and cannot be decorated, so smauglab/transforms/cpu/external.py builds their entries by hand. Two things there are declared per entry rather than assumed: wrap_random=False for transforms appended directly rather than inside a RandomTransform (they own no probability, so a config setting `p` is an error), and context_params for values nnU-Net supplies at runtime. The registered set and registry.PIPELINE_ORDER now agree exactly, 29 GPU and 22 CPU, and that is checked in both directions: registration rejects a class missing from the table, and a test rejects a table entry nothing registers. The `if` ladders still run and still decide what a config builds. Nothing reads the registry yet except the tests, which is deliberate -- this commit is only about the metadata being present and correct. Two things had to come with it: * RandomDomainTransferGPU still took **kwargs, so registering it failed the no-**kwargs-without-forwards_to rule. Removed. The previous commit's message claimed **kwargs was gone from every transform constructor; that was true of every file it touched, but it did not touch this one. * Its bank path was a hardcoded absolute path into one machine's NAS home directory, which is what external_asset exists to describe. It becomes the SMAUGLAB_DOMAIN_BANK environment variable, with no baked-in default and an error message that names the fix. Tests consult the registry entry rather than hardcoding the class, so a second such transform is covered automatically. Verified: with SMAUGLAB_DOMAIN_BANK pointing at the bank, all 24 seeded config digests are unchanged, including the domain-transfer config. Without it that one config skips rather than running, which is the only behaviour difference and the point of the change -- it was previously found by accident on exactly one machine. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/__init__.py | 24 +++- smauglab/transforms/cpu/artifact.py | 6 + smauglab/transforms/cpu/contrast.py | 53 +++++++++ smauglab/transforms/cpu/external.py | 72 ++++++++++++ smauglab/transforms/cpu/fromSeg.py | 7 ++ smauglab/transforms/cpu/spatial.py | 11 ++ smauglab/transforms/gpu/contrast.py | 101 +++++++++++++++++ smauglab/transforms/gpu/domain_transfer.py | 41 ++++++- smauglab/transforms/gpu/fromSeg.py | 11 ++ smauglab/transforms/gpu/spatial.py | 27 +++++ smauglab/transforms/synthseg/transforms.py | 13 ++- unit_tests/helpers.py | 35 +++++- unit_tests/test_registered_augmentations.py | 115 ++++++++++++++++++++ unit_tests/test_transforms_gpu.py | 24 +++- 14 files changed, 517 insertions(+), 23 deletions(-) create mode 100644 smauglab/transforms/cpu/external.py create mode 100644 unit_tests/test_registered_augmentations.py diff --git a/smauglab/transforms/__init__.py b/smauglab/transforms/__init__.py index 91a5776..f85ef14 100644 --- a/smauglab/transforms/__init__.py +++ b/smauglab/transforms/__init__.py @@ -1,9 +1,23 @@ """Augmentation transforms, split by execution backend. -`cpu` wraps batchgeneratorsv2 transforms for the dataloader worker; `gpu` wraps kornia -ones for the training step; `synthseg` holds the generative label-to-image augmentation. +`cpu` wraps batchgeneratorsv2 transforms for the dataloader worker; `gpu` wraps +kornia ones for the training step; `synthseg` holds the generative label-to-image +augmentation. -This is deliberately empty of imports for now. It becomes the point that populates -`smauglab.registry` -- importing it runs every `@register(...)` decorator -- once the -transform classes carry those decorators, which is the next change in this series. +Importing this package is what populates `smauglab.registry`: every augmentation +class carries an `@register(...)` decorator, so the registry is complete once these +modules have been imported and empty before. `registry.load_all()` does exactly +this import, which is why lookups are correct without callers having to know which +module defines what. + +Note this is deliberately NOT done from `smauglab/__init__.py`: it pulls in torch, +kornia and batchgeneratorsv2, and a bare `import smauglab` should not have to pay +for that. """ + +from smauglab.transforms.cpu import artifact, contrast, external, fromSeg, spatial # noqa: F401 +from smauglab.transforms.gpu import contrast as gpu_contrast # noqa: F401 +from smauglab.transforms.gpu import domain_transfer # noqa: F401 +from smauglab.transforms.gpu import fromSeg as gpu_fromSeg +from smauglab.transforms.gpu import spatial as gpu_spatial # noqa: F401 +from smauglab.transforms.synthseg import transforms as synthseg_transforms # noqa: F401 diff --git a/smauglab/transforms/cpu/artifact.py b/smauglab/transforms/cpu/artifact.py index 819bf02..6f6469d 100644 --- a/smauglab/transforms/cpu/artifact.py +++ b/smauglab/transforms/cpu/artifact.py @@ -2,6 +2,7 @@ import torchio as tio from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.cpu.torchio_ops import TransformFactory, apply_enabled, select #: Artifact name -> the torchio transform that produces it, in application order. @@ -17,6 +18,11 @@ } +@register( + aug_id=AugId.ARTIFACT, + backend=Backend.CPU, + group=AugType.TA, +) class ArtifactTransform(BasicTransform): def __init__(self, motion=False, ghosting=False, spike=False, bias_field=False, blur=False, noise=False, swap=False, random_pick=False): """ diff --git a/smauglab/transforms/cpu/contrast.py b/smauglab/transforms/cpu/contrast.py index f6b0fe9..896f13b 100644 --- a/smauglab/transforms/cpu/contrast.py +++ b/smauglab/transforms/cpu/contrast.py @@ -4,11 +4,19 @@ import torch.nn.functional as F from batchgeneratorsv2.helpers.scalar_type import RandomScalar from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform +from batchgeneratorsv2.transforms.intensity.contrast import BGContrast from batchgeneratorsv2.transforms.intensity.gamma import GammaTransform +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.kernels import laplace_kernel, scharr_kernels +@register( + aug_id=AugId.INV_GAMMA, + backend=Backend.CPU, + group=AugType.GE, + param_adapters={"gamma": BGContrast}, +) class InvertedGammaTransform(GammaTransform): """Gamma adjustment applied to the inverted image. @@ -99,6 +107,11 @@ def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: # either backend and `kernel_type` disappears from the config surface. +@register( + aug_id=AugId.LAPLACE, + backend=Backend.CPU, + group=AugType.TA, +) class LaplaceConvTransform(_ConvBaseTransform): """Laplacian edge enhancement.""" @@ -106,6 +119,11 @@ def __init__(self, absolute: bool = False, retain_stats: bool = False): super().__init__(kernel_type="Laplace", absolute=absolute, retain_stats=retain_stats) +@register( + aug_id=AugId.SCHARR, + backend=Backend.CPU, + group=AugType.TA, +) class ScharrConvTransform(_ConvBaseTransform): """Scharr gradient-magnitude edge filter.""" @@ -113,6 +131,11 @@ def __init__(self, absolute: bool = True, retain_stats: bool = False): super().__init__(kernel_type="Scharr", absolute=absolute, retain_stats=retain_stats) +@register( + aug_id=AugId.HISTOGRAM_EQUAL, + backend=Backend.CPU, + group=AugType.TA, +) class HistogramEqualTransform(ImageOnlyTransform): """ Update image intensity using histogram manipulations @@ -220,30 +243,55 @@ def __init__(self, retain_stats: bool = False): super().__init__(function=type(self).function_impl, retain_stats=retain_stats) +@register( + aug_id=AugId.FUNC_LOG1P, + backend=Backend.CPU, + group=AugType.TA, +) class Log1pTransform(_NamedFunctionTransform): """Apply log(1 + x).""" function_impl = staticmethod(_log1p) +@register( + aug_id=AugId.FUNC_SQRT, + backend=Backend.CPU, + group=AugType.TA, +) class SqrtTransform(_NamedFunctionTransform): """Apply sqrt(x).""" function_impl = staticmethod(torch.sqrt) +@register( + aug_id=AugId.FUNC_SIN, + backend=Backend.CPU, + group=AugType.TA, +) class SinTransform(_NamedFunctionTransform): """Apply sin(x).""" function_impl = staticmethod(torch.sin) +@register( + aug_id=AugId.FUNC_EXP, + backend=Backend.CPU, + group=AugType.TA, +) class ExpTransform(_NamedFunctionTransform): """Apply exp(x).""" function_impl = staticmethod(torch.exp) +@register( + aug_id=AugId.FUNC_SIGMOID, + backend=Backend.CPU, + group=AugType.TA, +) class SigmoidTransform(_NamedFunctionTransform): """Apply the logistic sigmoid 1 / (1 + exp(-x)).""" @@ -301,6 +349,11 @@ def apply_filter(x: torch.Tensor, kernel: torch.Tensor, **kwargs) -> torch.Tenso return output.view(batch, chns, *output.shape[2:]) +@register( + aug_id=AugId.ZSCORE, + backend=Backend.CPU, + group=AugType.GE, +) class ZscoreNormalization(ImageOnlyTransform): """ Z-score normalization of image diff --git a/smauglab/transforms/cpu/external.py b/smauglab/transforms/cpu/external.py new file mode 100644 index 0000000..8f5d0eb --- /dev/null +++ b/smauglab/transforms/cpu/external.py @@ -0,0 +1,72 @@ +"""Registry entries for the batchgeneratorsv2 transforms the CPU pipeline composes. + +These are third-party classes, so they cannot carry an `@register` decorator; the +entries are built here instead. Everything else about them is identical to a +decorated augmentation -- same config key rules, same signature-derived parameter +validation. + +Two things differ from the GPU side and are declared per entry: + +* `wrap_random=False` for transforms that are appended directly rather than inside + a `RandomTransform(...)`. Those own no application probability, so `p` is rejected + in a config for them. +* `context_params` for values nnU-Net supplies at runtime (patch size, rotation + range). A config must not set those, and the builder injects them. +""" + +from batchgeneratorsv2.transforms.intensity.brightness import MultiplicativeBrightnessTransform +from batchgeneratorsv2.transforms.intensity.contrast import BGContrast, ContrastTransform +from batchgeneratorsv2.transforms.intensity.gamma import GammaTransform +from batchgeneratorsv2.transforms.intensity.gaussian_noise import GaussianNoiseTransform +from batchgeneratorsv2.transforms.noise.gaussian_blur import GaussianBlurTransform +from batchgeneratorsv2.transforms.spatial.low_resolution import SimulateLowResolutionTransform +from batchgeneratorsv2.transforms.spatial.mirroring import MirrorTransform +from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform + +from smauglab.registry import AugEntry, AugId, AugType, Backend, register_entry + + +def _entry(cls: type, aug_id: AugId, group: AugType, **kwargs) -> AugEntry: + return register_entry( + AugEntry( + name=cls.__name__, + cls=cls, + backend=Backend.CPU, + aug_id=aug_id, + group=group, + summary=(cls.__doc__ or "").strip().split("\n", 1)[0], + **kwargs, + ) + ) + + +# Pipeline positions for these live in registry.PIPELINE_ORDER[Backend.CPU], taken +# from the sequence in AugTransforms._build_transforms. +_entry( + SpatialTransform, + AugId.SPATIAL, + AugType.GEO, + wrap_random=False, + # patch_size is positional and rotation comes from nnU-Net's + # configure_rotation_dummyDA_mirroring_and_inital_patch_size. + context_params=("patch_size", "rotation"), +) +_entry(GaussianNoiseTransform, AugId.GAUSSIAN_NOISE, AugType.GE) +_entry(GaussianBlurTransform, AugId.GAUSSIAN_BLUR, AugType.GE) +_entry( + MultiplicativeBrightnessTransform, + AugId.BRIGHTNESS, + AugType.GE, + param_adapters={"multiplier_range": BGContrast}, +) +_entry( + ContrastTransform, + AugId.CONTRAST, + AugType.GE, + param_adapters={"contrast_range": BGContrast}, +) +_entry(SimulateLowResolutionTransform, AugId.LOW_RES, AugType.GE) +_entry(GammaTransform, AugId.GAMMA, AugType.GE, param_adapters={"gamma": BGContrast}) +# Appended bare, and its allowed_axes comes from the config rather than the trainer: +# AugTransforms reads transform_params["mirror_axes"], not its own mirror_axes argument. +_entry(MirrorTransform, AugId.MIRROR, AugType.GEO, wrap_random=False) diff --git a/smauglab/transforms/cpu/fromSeg.py b/smauglab/transforms/cpu/fromSeg.py index 30e9f2c..d4e8b50 100644 --- a/smauglab/transforms/cpu/fromSeg.py +++ b/smauglab/transforms/cpu/fromSeg.py @@ -5,7 +5,14 @@ from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform from scipy.stats import norm +from smauglab.registry import AugId, AugType, Backend, register + +@register( + aug_id=AugId.REDISTRIBUTE_SEG, + backend=Backend.CPU, + group=AugType.TA, +) class RedistributeTransform(BasicTransform): """ Redistribute image values using segmentation regions. diff --git a/smauglab/transforms/cpu/spatial.py b/smauglab/transforms/cpu/spatial.py index 1511397..839f7e6 100644 --- a/smauglab/transforms/cpu/spatial.py +++ b/smauglab/transforms/cpu/spatial.py @@ -4,6 +4,7 @@ import torchio as tio from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform, ImageOnlyTransform +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.cpu.torchio_ops import TransformFactory, apply_enabled, select #: Transform name -> the torchio transform it runs, in application order. @@ -18,6 +19,11 @@ } +@register( + aug_id=AugId.SPATIAL_CUSTOM, + backend=Backend.CPU, + group=AugType.GEO, +) class SpatialCustomTransform(BasicTransform): def __init__(self, flip=False, affine=False, elastic=False, anisotropy=False, random_pick=False): """ @@ -48,6 +54,11 @@ def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tup ### Shape transform +@register( + aug_id=AugId.SHAPE, + backend=Backend.CPU, + group=AugType.GE, +) class ShapeTransform(ImageOnlyTransform): def __init__(self, shape_min=1, ignore_axes=()): """ diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index 1ad8f26..61ee817 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -7,6 +7,7 @@ from torch import Tensor from torch.nn import functional as F +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform from smauglab.transforms.kernels import gaussian_kernel3d, laplace_kernel, scharr_kernels from smauglab.transforms.rng import shared_choice @@ -331,6 +332,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # the ladder is gone. +@register( + aug_id=AugId.LAPLACE, + backend=Backend.GPU, + group=AugType.TA, +) class RandomLaplaceGPU(_RandomConvBaseGPU): """Laplacian edge enhancement.""" @@ -364,6 +370,11 @@ def __init__( ) +@register( + aug_id=AugId.SCHARR, + backend=Backend.GPU, + group=AugType.TA, +) class RandomScharrGPU(_RandomConvBaseGPU): """Scharr gradient-magnitude edge filter.""" @@ -397,6 +408,11 @@ def __init__( ) +@register( + aug_id=AugId.GAUSSIAN_BLUR, + backend=Backend.GPU, + group=AugType.GE, +) class RandomGaussianBlurGPU(_RandomConvBaseGPU): """Gaussian blur via separable convolution.""" @@ -430,6 +446,11 @@ def __init__( ) +@register( + aug_id=AugId.UNSHARP_MASK, + backend=Backend.GPU, + group=AugType.TA, +) class RandomUnsharpMaskGPU(_RandomConvBaseGPU): """Unsharp masking: sharpen by subtracting a blurred copy.""" @@ -465,6 +486,11 @@ def __init__( ) +@register( + aug_id=AugId.RAND_CONV, + backend=Backend.GPU, + group=AugType.TA, +) class RandomRandConvGPU(_RandomConvBaseGPU): """RandConv: convolution with a randomly drawn multi-scale kernel.""" @@ -535,6 +561,11 @@ def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torc ## Noise transform +@register( + aug_id=AugId.GAUSSIAN_NOISE, + backend=Backend.GPU, + group=AugType.GE, +) class RandomGaussianNoiseGPU(ImageOnlyTransform): """Add random Gaussian noise to image. If the image is torch Tensor, it is expected to have [N, C, X, Y] or [N, C, X, Y, Z] shape. @@ -597,6 +628,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ ## Multiplicative brightness transform +@register( + aug_id=AugId.BRIGHTNESS, + backend=Backend.GPU, + group=AugType.GE, +) class RandomBrightnessGPU(ImageOnlyTransform): """Apply random brightness adjustment to image. If the image is torch Tensor, it is expected to have [N, C, X, Y] or [N, C, X, Y, Z] shape. @@ -765,6 +801,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # cannot express the same augmentation two ways. +@register( + aug_id=AugId.GAMMA, + backend=Backend.GPU, + group=AugType.GE, +) class RandomGammaGPU(_RandomGammaBaseGPU): """Random gamma adjustment.""" @@ -796,6 +837,11 @@ def __init__( ) +@register( + aug_id=AugId.INV_GAMMA, + backend=Backend.GPU, + group=AugType.GE, +) class RandomInvGammaGPU(_RandomGammaBaseGPU): """Random gamma adjustment applied to the inverted image.""" @@ -828,6 +874,11 @@ def __init__( ## nnunetv2 contrast transform +@register( + aug_id=AugId.CONTRAST, + backend=Backend.GPU, + group=AugType.GE, +) class RandomContrastGPU(ImageOnlyTransform): """Apply random gamma adjustment to image. If the image is torch Tensor, it is expected to have [N, C, X, Y] or [N, C, X, Y, Z] shape. @@ -1031,30 +1082,55 @@ def __init__( ) +@register( + aug_id=AugId.FUNC_LOG1P, + backend=Backend.GPU, + group=AugType.TA, +) class RandomLog1pGPU(_RandomNamedFunctionGPU): """Apply log(1 + x).""" function = staticmethod(_log1p) +@register( + aug_id=AugId.FUNC_SQRT, + backend=Backend.GPU, + group=AugType.TA, +) class RandomSqrtGPU(_RandomNamedFunctionGPU): """Apply sqrt(x).""" function = staticmethod(torch.sqrt) +@register( + aug_id=AugId.FUNC_SIN, + backend=Backend.GPU, + group=AugType.TA, +) class RandomSinGPU(_RandomNamedFunctionGPU): """Apply sin(x).""" function = staticmethod(torch.sin) +@register( + aug_id=AugId.FUNC_EXP, + backend=Backend.GPU, + group=AugType.TA, +) class RandomExpGPU(_RandomNamedFunctionGPU): """Apply exp(x).""" function = staticmethod(torch.exp) +@register( + aug_id=AugId.FUNC_SIGMOID, + backend=Backend.GPU, + group=AugType.TA, +) class RandomSigmoidGPU(_RandomNamedFunctionGPU): """Apply the logistic sigmoid 1 / (1 + exp(-x)).""" @@ -1062,6 +1138,11 @@ class RandomSigmoidGPU(_RandomNamedFunctionGPU): ## Inverse transform +@register( + aug_id=AugId.INVERSE, + backend=Backend.GPU, + group=AugType.TA, +) class RandomInverseGPU(ImageOnlyTransform): """Inverse image based on probability. If the image is torch Tensor, it is expected to have [N, C, X, Y] or [N, C, X, Y, Z] shape. @@ -1133,6 +1214,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ ## Histogram transform +@register( + aug_id=AugId.HISTOGRAM_EQUAL, + backend=Backend.GPU, + group=AugType.TA, +) class RandomHistogramEqualizationGPU(ImageOnlyTransform): """Apply histogram equalization transformation to the image based on probability. If the image is torch Tensor, it is expected to have [N, C, X, Y] or [N, C, X, Y, Z] shape. @@ -1227,6 +1313,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ return input +@register( + aug_id=AugId.BIAS_FIELD, + backend=Backend.GPU, + group=AugType.TA, +) class RandomBiasFieldGPU(ImageOnlyTransform): """Apply a smooth multiplicative bias field to selected channels. @@ -1408,6 +1499,11 @@ def apply_transform( # Random clamping transform +@register( + aug_id=AugId.CLAMP, + backend=Backend.GPU, + group=AugType.GE, +) class RandomClampGPU(ImageOnlyTransform): """Apply random gamma adjustment to image. If the image is torch Tensor, it is expected to have [N, C, X, Y] or [N, C, X, Y, Z] shape. @@ -1483,6 +1579,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ return input +@register( + aug_id=AugId.ZSCORE, + backend=Backend.GPU, + group=AugType.GE, +) class ZscoreNormalizationGPU(ImageOnlyTransform): """Apply z-score normalization to selected channels. diff --git a/smauglab/transforms/gpu/domain_transfer.py b/smauglab/transforms/gpu/domain_transfer.py index 32414e8..e46a193 100644 --- a/smauglab/transforms/gpu/domain_transfer.py +++ b/smauglab/transforms/gpu/domain_transfer.py @@ -35,6 +35,7 @@ """ import math +import os from typing import Any import numpy as np @@ -43,11 +44,35 @@ from torch.distributions import Dirichlet from torch.nn import functional as F +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform from smauglab.transforms.kernels import gaussian_blur3d, random_bias_field3d # Default transfer LUT bank (built by embeddaug/analysis/playground/build_transfer_bank.py). -DEFAULT_BANK_PATH = "/DATA/NAS/ongoing_projects/hendrik/nathan-transferaug/embeddaug/analysis/playground/results/domain_transfer_bank.npz" +# The transfer LUT bank is a multi-hundred-MB artefact built offline by +# embeddaug/analysis/playground/build_transfer_bank.py, so it is not shipped in the +# wheel. Point this env var at it; there is deliberately no baked-in default, because +# the previous one was an absolute path into a single machine's NAS home directory and +# silently made this transform unusable for everyone else. +BANK_PATH_ENV_VAR = "SMAUGLAB_DOMAIN_BANK" + + +def resolve_bank_path(bank_path: str | None = None) -> str: + """Locate the domain-transfer LUT bank, explicit argument first, then the env var. + + Raises with the fix spelled out rather than letting np.load report a bare + FileNotFoundError on a path the caller never chose. + """ + resolved = bank_path or os.environ.get(BANK_PATH_ENV_VAR) + if not resolved: + raise FileNotFoundError( + "RandomDomainTransferGPU needs a transfer LUT bank. Pass bank_path=..., set it in " + f"the config, or export {BANK_PATH_ENV_VAR}=/path/to/domain_transfer_bank.npz " + "(built by embeddaug/analysis/playground/build_transfer_bank.py)." + ) + if not os.path.isfile(resolved): + raise FileNotFoundError(f"Domain transfer bank not found at {resolved!r} (from {BANK_PATH_ENV_VAR} or bank_path).") + return resolved def _gaussian_blur3d(x: torch.Tensor, sigma: float) -> torch.Tensor: @@ -91,6 +116,13 @@ def _random_smooth_field01(shape, scale: float, gain: float, device, dtype) -> t return torch.sigmoid(gain * field + offset) +@register( + aug_id=AugId.DOMAIN_TRANSFER, + backend=Backend.GPU, + group=AugType.TA, + external_asset=BANK_PATH_ENV_VAR, + smoke_kwargs={"any_source": True}, +) class RandomDomainTransferGPU(ImageOnlyTransform): """Randomly transfer an image's appearance to another sequence/cluster (see module docstring).""" @@ -115,14 +147,13 @@ def __init__( spatial_mix_gain: float = 3.0, same_on_batch: bool = False, p: float = 0.2, + p_batch: float = 1.0, keepdim: bool = True, - **kwargs, ) -> None: if apply_to_channel is None: apply_to_channel = [0] - super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) - bank_path = bank_path or DEFAULT_BANK_PATH - data = np.load(bank_path, allow_pickle=True) + super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) + data = np.load(resolve_bank_path(bank_path), allow_pickle=True) self.labels: list[str] = [str(x) for x in data["labels"].tolist()] self.L = int(data["L"]) self.num_classes = int(data["num_classes"]) diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index 71ed118..4da6340 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -5,6 +5,7 @@ from torch import Tensor, nn from torch.nn import functional as F +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform from smauglab.transforms.kernels import gaussian_blur3d from smauglab.transforms.rng import shared_choice @@ -89,6 +90,11 @@ def _normal_pdf(x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch ## Redistribute segmentation values transform (GPU) +@register( + aug_id=AugId.REDISTRIBUTE_SEG, + backend=Backend.GPU, + group=AugType.TA, +) class RandomRedistributeSegGPU(ImageOnlyTransform): """Redistribute image values using segmentation regions (GPU version). @@ -258,6 +264,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ return input +@register( + aug_id=AugId.PALETTE, + backend=Backend.GPU, + group=AugType.TA, +) class RandomPaletteGPU(ImageOnlyTransform): """ SmaugLab GPU augmentation implementing PALETTE synthesis. diff --git a/smauglab/transforms/gpu/spatial.py b/smauglab/transforms/gpu/spatial.py index 9d3cbd6..f46b408 100644 --- a/smauglab/transforms/gpu/spatial.py +++ b/smauglab/transforms/gpu/spatial.py @@ -18,10 +18,16 @@ # the kornia-compat matrix in tests.yml exercises, so it stays. from kornia.core.utils import _extract_device_dtype # type: ignore[no-redef] +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform # Affine transform +@register( + aug_id=AugId.AFFINE, + backend=Backend.GPU, + group=AugType.GEO, +) class RandomAffineGPU(RigidAffineAugmentationBase3D): r"""Apply affine transformation 3D volumes (5D tensor). @@ -200,6 +206,12 @@ def apply_transform_mask( # Low resolution transform +@register( + aug_id=AugId.LOW_RES, + backend=Backend.GPU, + group=AugType.GE, + force_sequential=True, +) class RandomLowResTransformGPU(RigidAffineAugmentationBase3D): """ Apply low resolution simulation to 3D volumes (5D tensor). @@ -348,6 +360,11 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> # Acquisition transforms +@register( + aug_id=AugId.ACQ, + backend=Backend.GPU, + group=AugType.GE, +) class RandomAcqTransformGPU(ImageOnlyTransform): """ Randomly lower acquisition along one axes only. @@ -434,6 +451,11 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # Flip transforms +@register( + aug_id=AugId.FLIP, + backend=Backend.GPU, + group=AugType.GEO, +) class RandomFlipTransformGPU(RigidAffineAugmentationBase3D): """ Apply low resolution simulation to 3D volumes (5D tensor). @@ -564,6 +586,11 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> # Crop transform +@register( + aug_id=AugId.CROP, + backend=Backend.GPU, + group=AugType.GEO, +) class RandomCropTransformGPU(RigidAffineAugmentationBase3D): """ Apply low resolution simulation to 3D volumes (5D tensor). diff --git a/smauglab/transforms/synthseg/transforms.py b/smauglab/transforms/synthseg/transforms.py index af8f381..2017ef1 100644 --- a/smauglab/transforms/synthseg/transforms.py +++ b/smauglab/transforms/synthseg/transforms.py @@ -28,6 +28,7 @@ import torch from torch import Tensor, nn +from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform from smauglab.transforms.synthseg.generator import SynthSegGenerator @@ -82,6 +83,15 @@ def _filter_generator_kwargs(params: dict[str, Any]) -> dict[str, Any]: return {k: v for k, v in params.items() if k in _GENERATOR_KEYS} +@register( + aug_id=AugId.SYNTHSEG, + backend=Backend.GPU, + group=AugType.TA, + forwards_to=SynthSegGenerator, + # Forced below to keep the synthesis intensity-only; a config setting any of + # these would be silently overridden, so they are rejected instead. + context_params=("apply_affine", "apply_nonlinear", "flipping", "output_shape"), +) class RandomSynthSegGPU(ImageOnlyTransform): """Replace the image with a SynthSeg GMM synthesis of ``params['seg']``. @@ -104,10 +114,11 @@ def __init__( apply_to_channel: list[int] | None = None, same_on_batch: bool = False, p: float = 0.5, + p_batch: float = 1.0, keepdim: bool = True, **kwargs: Any, ) -> None: - super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) + 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] gen_kwargs = _filter_generator_kwargs(kwargs) # Intensity-only: never deform/flip internally (geometry comes from the diff --git a/unit_tests/helpers.py b/unit_tests/helpers.py index 53aacf2..360d20e 100644 --- a/unit_tests/helpers.py +++ b/unit_tests/helpers.py @@ -90,17 +90,40 @@ def gpu_config_paths() -> list[Path]: def requires_external_asset(config_path: Path) -> str | None: """Return a skip reason if a config needs an asset that is not on this machine. - RandomDomainTransferGPU loads a precomputed histogram bank from an absolute - path baked into the module, which only exists on the authors' machines. - Rather than fail CI, skip those configs and say why. + RandomDomainTransferGPU loads a precomputed histogram bank that is far too large + to ship in the wheel; its location comes from an environment variable. Rather than + fail CI, skip those configs and say why. """ - from smauglab.transforms.gpu.domain_transfer import DEFAULT_BANK_PATH + import os + + from smauglab.transforms.gpu.domain_transfer import BANK_PATH_ENV_VAR params = json.loads(config_path.read_text()) params = params.get("GPU", params) if not isinstance(params, dict): return None uses_transfer = params.get("RandomDomainTransferGPU") or params.get("DomainTransferTransform") - if uses_transfer and not Path(DEFAULT_BANK_PATH).is_file(): - return f"domain transfer bank not available at {DEFAULT_BANK_PATH}" + if not uses_transfer: + return None + configured = uses_transfer.get("bank_path") if isinstance(uses_transfer, dict) else None + bank = configured or os.environ.get(BANK_PATH_ENV_VAR) + if not bank or not Path(bank).is_file(): + return f"domain transfer bank not available (set ${BANK_PATH_ENV_VAR})" return None + + +def domain_bank_missing(entry) -> str | None: + """Skip reason if a registry entry needs an external artefact that is absent. + + `external_asset` names the environment variable that points at it. The + domain-transfer LUT bank is the only such artefact: hundreds of megabytes, built + offline, deliberately not shipped in the wheel. + """ + import os + + if not getattr(entry, "external_asset", None): + return None + location = os.environ.get(entry.external_asset) + if location and Path(location).is_file(): + return None + return f"{entry.name} needs ${entry.external_asset} to point at its data" diff --git a/unit_tests/test_registered_augmentations.py b/unit_tests/test_registered_augmentations.py new file mode 100644 index 0000000..0a4bf9f --- /dev/null +++ b/unit_tests/test_registered_augmentations.py @@ -0,0 +1,115 @@ +"""The real registry, and the artifacts generated from it. + +test_registry.py covers the mechanism against synthetic classes. This file covers +the actual augmentations: that every one is registered coherently, that the +generated matrix and template config are not stale, and that every registered +augmentation is genuinely constructible. + +Together these are the answer to "which augmentations exist and which backends +have them" -- previously recoverable only by reading four `if` ladders side by side. +""" + +from __future__ import annotations + +import unittest +from pathlib import Path + +from smauglab import registry +from smauglab.registry import AugId, Backend + +REPO = Path(__file__).resolve().parent.parent +TEMPLATE = REPO / "smauglab" / "configs" / "all_augmentations.json" + + +class TestRegistryIsPopulated(unittest.TestCase): + def test_both_backends_have_augmentations(self): + self.assertGreaterEqual(len(registry.names(Backend.GPU)), 25) + self.assertGreaterEqual(len(registry.names(Backend.CPU)), 20) + + def test_no_monai_implementations_yet(self): + """Tracked, not built. If this starts failing, the matrix gained a real cell.""" + self.assertEqual(registry.names(Backend.MONAI), []) + + def test_every_aug_id_is_used(self): + """An unused AugId is a concept nothing implements -- almost always a typo.""" + used = {entry.aug_id for entry in registry.entries()} + self.assertEqual(set(AugId) - used, set(), "AugId members with no implementation on any backend") + + def test_class_name_is_the_config_key(self): + for entry in registry.entries(): + with self.subTest(entry=entry.name): + self.assertEqual(entry.cls.__name__, entry.name) + + def test_every_registered_class_has_a_pipeline_position(self): + """PIPELINE_ORDER and the decorators must describe the same set, in both + directions: registration rejects an unlisted class, and this catches a listed + name that nothing registers (a typo, or a class that was renamed).""" + for backend in (Backend.GPU, Backend.CPU): + with self.subTest(backend=backend.value): + registered = {entry.name for entry in registry.entries(backend)} + listed = set(registry.PIPELINE_ORDER[backend]) + self.assertEqual(listed, registered) + + def test_no_registered_class_hides_parameters_behind_kwargs(self): + """**kwargs would make signature-derived validation accept anything.""" + for entry in registry.entries(): + with self.subTest(entry=entry.name): + if registry._has_var_keyword(entry.cls): + self.assertIsNotNone( + entry.forwards_to, + f"{entry.name} takes **kwargs without declaring forwards_to", + ) + + def test_gpu_transforms_expose_the_kornia_probability_parameters(self): + for entry in registry.entries(Backend.GPU): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + self.assertIn("p", accepted) + self.assertIn("p_batch", accepted) + + def test_probability_is_never_a_parameter_name(self): + """It was renamed to `p`; a survivor would mean a half-done migration.""" + for entry in registry.entries(): + with self.subTest(entry=entry.name): + self.assertNotIn("probability", registry.accepted_params(entry)) + + def test_legacy_config_keys_do_not_resolve(self): + """The hard break, on the names that actually appear in the old configs.""" + for legacy in ( + "ScharrTransform", + "UnsharpMaskTransform", + "RandomConvTransform", + "SynthSeg", + "AffineTransform", + "FlipTransform", + "RandomPALETTETransform", + "GammaTransform_invert", + "ImageContrastGPUTransform", + "PaletteSynthesisTransform", + ): + with self.subTest(legacy=legacy), self.assertRaises(registry.UnknownAugmentationError): + registry.get(legacy, Backend.GPU) + + +class TestEveryEntryIsConstructible(unittest.TestCase): + """The registry may not advertise an augmentation that cannot be built.""" + + def test_constructible_with_declared_defaults(self): + for entry in registry.entries(): + with self.subTest(entry=f"{entry.backend.value}.{entry.name}"): + if entry.external_asset: + from unit_tests.helpers import domain_bank_missing + + reason = domain_bank_missing(entry) + if reason: + self.skipTest(reason) + required = registry.required_params(entry) + if required: + # Legitimate: a few third-party CPU transforms take mandatory + # arguments that only a config or the trainer can supply. + self.assertTrue( + entry.backend is Backend.CPU or entry.context_params, + f"{entry.name} requires {sorted(required)} but nothing supplies them", + ) + continue + entry.cls(**dict(entry.smoke_kwargs)) diff --git a/unit_tests/test_transforms_gpu.py b/unit_tests/test_transforms_gpu.py index 9a5a0b1..b466a04 100644 --- a/unit_tests/test_transforms_gpu.py +++ b/unit_tests/test_transforms_gpu.py @@ -12,6 +12,7 @@ import importlib import inspect +import os import unittest from pathlib import Path @@ -76,13 +77,24 @@ def build_kwargs(cls, signature) -> dict: def skip_reason(cls) -> str | None: - """Some transforms depend on assets that do not exist on a fresh checkout.""" - if cls.__name__ == "RandomDomainTransferGPU": - from smauglab.transforms.gpu.domain_transfer import DEFAULT_BANK_PATH + """Some transforms depend on assets that do not exist on a fresh checkout. - if not Path(DEFAULT_BANK_PATH).is_file(): - return f"domain transfer bank not available at {DEFAULT_BANK_PATH}" - return None + Which ones is registry data now (`external_asset` names the environment variable + that points at the artefact), rather than a hardcoded class name here -- so a + second such transform is covered the moment it is registered. + """ + from smauglab import registry + + try: + entry = registry.get(cls.__name__) + except registry.UnknownAugmentationError: + return None + if not entry.external_asset: + return None + location = os.environ.get(entry.external_asset) + if location and Path(location).is_file(): + return None + return f"{cls.__name__} needs ${entry.external_asset} to point at its data" class TestTransformDiscovery(unittest.TestCase):