feat: register every augmentation

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

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):
         """

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,20 +107,35 @@ 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."""
 
     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."""
 
     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

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)

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.

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=()):
         """