From 6e9e156b99e6ad986f1777804f515e26c1254b7e Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 13:07:36 +0000 Subject: [PATCH] feat: add the augmentation registry smauglab/registry.py, with no consumers yet -- reviewable on its own before anything depends on it. It answers four questions that are currently answered by reading an `if` ladder: what augmentations exist per backend, which class a config key maps to, what parameters that class accepts, and which backends implement a given concept. Accepted parameters come from inspect.signature, so there is no second schema to drift out of step with the constructors; registering a class whose __init__ ends in **kwargs without declaring forwards_to is an error, because signature-based validation would otherwise accept anything. The module is stdlib-only and imports nothing from smauglab.transforms at module scope. That keeps the import graph acyclic -- transform modules import register from here -- and lets a CLI answer "what exists?" without paying for torch. On pipeline order. Each augmentation's position is data the ladders held implicitly, so the registry has to hold it explicitly; config key order cannot decide it, because the two have never agreed and honouring the file would silently reorder every pipeline the moment someone tidied a config. It is one PIPELINE_ORDER tuple per backend rather than an `order=` integer on each @register. The question a reviewer needs to answer is "does the new pipeline run things in the same sequence as the old one?", and against a single list that is a diff against the ladder it was derived from. Against integers spread over eight files it is not. Each GPU entry carries the ladder key it came from as a comment, so the correspondence is checkable line by line. Registering a class absent from its backend's tuple is an error, so the table cannot fall out of date. RandomLaplaceGPU is the one GPU entry with no ladder key: the GPU Laplace kernel was reachable only by setting kernel_type="Laplace" on ScharrTransform, so its position is a new choice, placed next to its sibling. isolated() gains an `order` argument. Without one it turns the position check off for the block, which is what most tests want -- they register throwaway classes whose pipeline position is not the thing under test. With one, registration is checked against it, which is how the check itself is covered. Also: smauglab stops being a PEP 420 namespace package. Every directory gets an __init__.py, because the registry needs a deterministic import-time population point and because a namespace package lets a stray smauglab/ elsewhere on sys.path silently merge into this one. namespace_packages/explicit_package_bases come out of [tool.mypy] accordingly. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 12 +- smauglab/__init__.py | 18 + smauglab/registry.py | 665 ++++++++++++++++++++++++++++ smauglab/transforms/__init__.py | 9 + smauglab/transforms/cpu/__init__.py | 1 + smauglab/transforms/gpu/__init__.py | 1 + smauglab/utils/__init__.py | 13 + unit_tests/test_registry.py | 260 +++++++++++ 8 files changed, 972 insertions(+), 7 deletions(-) create mode 100644 smauglab/__init__.py create mode 100644 smauglab/registry.py create mode 100644 smauglab/transforms/__init__.py create mode 100644 smauglab/transforms/cpu/__init__.py create mode 100644 smauglab/transforms/gpu/__init__.py create mode 100644 smauglab/utils/__init__.py create mode 100644 unit_tests/test_registry.py diff --git a/pyproject.toml b/pyproject.toml index e2ae4b1..a00694a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,9 +41,11 @@ keywords = [ "offline augmentation", "mri", ] -# smauglab has no __init__.py anywhere, so it is a PEP 420 namespace package. -# poetry-core still walks the tree correctly; the CI build job and -# unit_tests/test_packaging.py assert the wheel really contains every module. +# smauglab is a regular package: every directory has an __init__.py. It used to be a +# PEP 420 namespace package, which was dropped because the augmentation registry needs +# a deterministic import-time population point, and because a namespace package lets a +# stray smauglab/ elsewhere on sys.path silently merge into this one. +# The CI build job and unit_tests/test_packaging.py assert the wheel contains every module. packages = [{ include = "smauglab" }] # The config JSONs are package data, not code, so they need listing explicitly. # A wheel without them is broken: smauglab resolves its default config through @@ -258,10 +260,6 @@ python_version = "3.10" # incrementally rather than in one pass. ignore_missing_imports = true warn_redundant_casts = true -# smauglab/, smauglab/transforms/ and smauglab/utils/ have no __init__.py (PEP 420, -# see the `packages` note above), so mypy cannot derive module names without these. -namespace_packages = true -explicit_package_bases = true # Only the shipped package is gated. Tests and standalone scripts are not installed # and assert against literals, which reads badly to a type checker. exclude = ["unit_tests/", "scripts/"] diff --git a/smauglab/__init__.py b/smauglab/__init__.py new file mode 100644 index 0000000..19fee45 --- /dev/null +++ b/smauglab/__init__.py @@ -0,0 +1,18 @@ +"""SmaugLab -- data augmentation strategies for MRI segmentation training. + +Deliberately kept free of imports from `smauglab.transforms`. Pulling the +transforms in here would make a bare `import smauglab` drag in torch, kornia and +batchgeneratorsv2 (several seconds), which every console-script invocation would +then pay for. Import the subpackage you actually need: + + from smauglab.transforms.gpu.transforms import AugTransformsGPU +""" + +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("smauglab") +except PackageNotFoundError: # running from a source tree that was never installed + __version__ = "0.0.0" + +__all__ = ["__version__"] diff --git a/smauglab/registry.py b/smauglab/registry.py new file mode 100644 index 0000000..35b4887 --- /dev/null +++ b/smauglab/registry.py @@ -0,0 +1,665 @@ +"""The single source of truth for which augmentations exist. + +Each augmentation class registers itself with `@register(...)`. The registry then +answers: + +* what exists, per backend -- `names()`, `entries()` +* which class a config key maps to -- `get()` +* what parameters that class accepts -- `accepted_params()` +* which backends implement a given concept -- `matrix()`, keyed on `AugId` + +Deliberately stdlib-only, and deliberately importing nothing from +`smauglab.transforms` at module scope: that is what keeps the import graph acyclic +(transform modules import `register` from here) and lets a CLI ask "what exists?" +without paying for torch until it actually has to. +""" + +from __future__ import annotations + +import contextlib +import difflib +import importlib +import inspect +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, TypeVar + +__all__ = [ + "PIPELINE_ORDER", + "AugEntry", + "AugId", + "AugType", + "Backend", + "InvalidConfigError", + "RegistryError", + "UnknownAugmentationError", + "UnknownParameterError", + "accepted_params", + "entries", + "get", + "isolated", + "matrix", + "names", + "pipeline_position", + "register", + "register_entry", +] + + +# `StrEnum` is 3.11+; pyproject promises `requires-python = ">=3.10"`. +class Backend(str, Enum): + """Where an augmentation runs, which is also which config section it lives in.""" + + GPU = "GPU" + CPU = "CPU" + MONAI = "MONAI" + + +class AugType(str, Enum): + """Augmentation strength class. + + Member names are verbatim from segtransferaug's `AUG2GROUP`, which this + replaces, so downstream `AugType.TA` keeps resolving. The GPU random-order + pipelines bucket transforms by this to decide what goes inside each + `RandomChooseXTransformsGPU`. + """ + + GEO = "GEO" # geometric: applied in sequence, never bucketed + GE = "GE" # general enhancement, usually light + TA = "TA" # transfer augmentation, usually heavy + + +class AugId(str, Enum): + """A backend-neutral augmentation *concept*. + + Config keys are class names and therefore backend-specific + (`RandomGaussianNoiseGPU` vs batchgeneratorsv2's `GaussianNoiseTransform`), so + a join key is needed to say "these are the same augmentation on two backends". + That is what makes `matrix()` -- the CPU/GPU/MONAI coverage table -- possible. + + An enum rather than a free string on purpose: a typo'd `"gausian_noise"` would + silently orphan a matrix row, whereas `AugId.GAUSIAN_NOISE` fails at import. + The member list doubles as the authoritative inventory of concepts, which is + what "track which augmentations have a MONAI version" actually means -- the row + exists, and an empty cell is the record that no implementation does. + """ + + # -- geometric + FLIP = "flip" + AFFINE = "affine" + CROP = "crop" + SPATIAL = "spatial" # nnU-Net's own SpatialTransform + # -- general enhancement + GAUSSIAN_NOISE = "gaussian_noise" + GAUSSIAN_BLUR = "gaussian_blur" + BRIGHTNESS = "brightness" + CONTRAST = "contrast" + GAMMA = "gamma" + INV_GAMMA = "inv_gamma" + CLAMP = "clamp" + LOW_RES = "low_res" + ACQ = "acq" + ZSCORE = "zscore" + MIRROR = "mirror" + # -- transfer augmentation + SCHARR = "scharr" + LAPLACE = "laplace" + UNSHARP_MASK = "unsharp_mask" + RAND_CONV = "rand_conv" + BIAS_FIELD = "bias_field" + INVERSE = "inverse" + HISTOGRAM_EQUAL = "histogram_equal" + REDISTRIBUTE_SEG = "redistribute_seg" + PALETTE = "palette" + DOMAIN_TRANSFER = "domain_transfer" + SYNTHSEG = "synthseg" + ARTIFACT = "artifact" + SPATIAL_CUSTOM = "spatial_custom" + SHAPE = "shape" + # -- elementwise functions, one concept each so key <-> class stays 1:1 + FUNC_LOG1P = "func_log1p" + FUNC_SQRT = "func_sqrt" + FUNC_SIN = "func_sin" + FUNC_EXP = "func_exp" + FUNC_SIGMOID = "func_sigmoid" + + +class RegistryError(Exception): + """Base class for every registry and config-resolution failure.""" + + +class UnknownAugmentationError(RegistryError, KeyError): + """A config named an augmentation that is not registered for that backend.""" + + def __str__(self) -> str: # KeyError.__str__ would repr() the message + return self.args[0] if self.args else "" + + +class UnknownParameterError(RegistryError, TypeError): + """A config passed a parameter the augmentation's constructor does not accept.""" + + +class InvalidConfigError(RegistryError, ValueError): + """A config has one or more problems. Every problem is reported at once, so a + broken file is fixed in one pass rather than one pytest run at a time.""" + + def __init__(self, source: str, problems: list[str]) -> None: + self.source = source + self.problems = problems + joined = "\n".join(f" - {p}" for p in problems) + super().__init__(f"{source}: {len(problems)} problem(s)\n{joined}") + + +# Renamed *parameters*, for diagnostics only -- consulted when building an error +# message, NEVER when loading a config. difflib cannot bridge these on its own: +# "probability" vs "p" scores ~0.17, well under any usable cutoff, so without this +# table the single most common migration mistake would get no suggestion at all. +# A test asserts that no key here resolves through `get()`, which is what keeps it +# from becoming a back-compat path. +# +# Renamed augmentations are deliberately absent: stem matching in +# `_unknown_augmentation_message` already bridges "ScharrTransform" -> RandomScharrGPU +# and "SynthSeg" -> RandomSynthSegGPU without a table to maintain. +RENAMED_HINTS: Mapping[str, str] = MappingProxyType( + { + "probability": "p", + "shear": "shears", + "invert_image": "use RandomInvGammaGPU instead", + "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", + } +) + + +#: Where each augmentation sits in its backend's pipeline. +#: +#: This replaces the fixed sequence that used to be written into the `if` ladders -- +#: `gpu/transforms.py` for GPU, `cpu/transforms.py` for CPU. It is one list per backend +#: rather than an `order=` number on each `@register` for a specific reason: the whole +#: point of review here is "does the new pipeline run things in the same sequence as the +#: old one?", and that question is answerable by diffing this against the ladder it came +#: from. Spread over eight files as integers, it is not. +#: +#: Order matters. Config key order does not decide it -- the two have never agreed, and +#: honouring the file would silently reorder every pipeline the moment someone tidied a +#: config. A config can opt into key order explicitly; see `pipeline.order` in +#: smauglab/config.py. +#: +#: Registering a class that is absent from its backend's tuple is an error, so this +#: cannot silently fall out of date. +PIPELINE_ORDER: Mapping[Backend, tuple[str, ...]] = MappingProxyType( + { + # From the ladder in gpu/transforms.py, top to bottom. The four names that ladder + # reached through a `kernel_type` / `func` / `invert_image` argument are expanded + # into the leaf classes that replaced them, in the position the argument was read. + Backend.GPU: ( + "RandomFlipTransformGPU", # FlipTransform + "RandomAffineGPU", # AffineTransform + "RandomSynthSegGPU", # SynthSeg + "RandomPaletteGPU", # RandomPALETTETransform + "RandomDomainTransferGPU", # DomainTransferTransform + "RandomInverseGPU", # InverseTransform + "RandomHistogramEqualizationGPU", # HistogramEqualizationTransform + "RandomRedistributeSegGPU", # RedistributeSegTransform + "RandomScharrGPU", # ScharrTransform + # No ladder key: the GPU Laplace kernel was reachable only by setting + # kernel_type="Laplace" on ScharrTransform. Placed next to its sibling. + "RandomLaplaceGPU", + "RandomUnsharpMaskGPU", # UnsharpMaskTransform + "RandomRandConvGPU", # RandomConvTransform + "RandomClampGPU", # ClampTransform + "RandomGaussianNoiseGPU", # GaussianNoiseTransform + "RandomGaussianBlurGPU", # GaussianBlurTransform + "RandomBrightnessGPU", # BrightnessTransform + "RandomGammaGPU", # GammaTransform + "RandomInvGammaGPU", # InvGammaTransform + "RandomContrastGPU", # ContrastTransform + # FunctionTransform built one transform per entry in its `func` list, in this + # order; the leaves keep it. + "RandomLog1pGPU", + "RandomSqrtGPU", + "RandomSinGPU", + "RandomExpGPU", + "RandomSigmoidGPU", + "RandomLowResTransformGPU", # SimulateLowResTransform + "RandomAcqTransformGPU", # AcqTransform + "RandomCropTransformGPU", # CropTransform + "RandomBiasFieldGPU", # BiasFieldTransform + "ZscoreNormalizationGPU", # ZscoreNormalizationTransform + ), + # From the ladder in cpu/transforms.py, top to bottom. + Backend.CPU: ( + "LaplaceConvTransform", + "ScharrConvTransform", + "Log1pTransform", + "SqrtTransform", + "SinTransform", + "ExpTransform", + "SigmoidTransform", + "HistogramEqualTransform", + "RedistributeTransform", + "ShapeTransform", + "ArtifactTransform", + "SpatialCustomTransform", + "SpatialTransform", + "GaussianNoiseTransform", + "GaussianBlurTransform", + "MultiplicativeBrightnessTransform", + "ContrastTransform", + "SimulateLowResolutionTransform", + "InvertedGammaTransform", + "GammaTransform", + "MirrorTransform", + "ZscoreNormalization", + ), + # No MONAI augmentations are implemented yet; the empty tuple is the record of + # that, and is what `matrix()` renders as an empty column. + Backend.MONAI: (), + } +) + + +def pipeline_position(entry: AugEntry) -> int: + """Index of `entry` in its backend's pipeline order. + + Unlisted names sort last rather than raising: only `isolated()` can produce one, + and a test's throwaway class should not make `entries()` explode. + """ + order = PIPELINE_ORDER[entry.backend] + return order.index(entry.name) if entry.name in order else len(order) + + +@dataclass(frozen=True, slots=True) +class AugEntry: + """One registered augmentation: a class plus everything a config needs to know.""" + + name: str + cls: type + backend: Backend + aug_id: AugId + group: AugType + # Set when the constructor legitimately forwards **kwargs to another class, so + # `accepted_params` unions both signatures instead of giving up. RandomSynthSegGPU + # forwards ~38 parameters to SynthSegGenerator. + forwards_to: type | None = None + # Supplied by the builder at runtime (nnU-Net hands over patch size, rotation and + # mirror axes), so these are rejected if a config tries to set them. + context_params: tuple[str, ...] = () + # CPU only: batchgeneratorsv2 puts the apply probability on a RandomTransform + # wrapper rather than the transform, so `p` is a builder key, not a ctor kwarg. + wrap_random: bool = True + # Run in pipeline order even in the random-order pipelines, instead of going + # into a shuffled RandomChooseX bucket. GEO transforms are always sequential; + # this is for the ones whose group says otherwise. RandomLowResTransformGPU is + # the only case: segtransferaug's AUG2GROUP calls it GE, but the random-order + # builder has always run it in sequence, and the group is used downstream for + # filtering, so the two meanings are kept apart rather than reconciled. + force_sequential: bool = False + # Name of an env var pointing at a large artefact the transform needs but the + # wheel does not ship. Tests skip rather than fail when it is unset. + external_asset: str | None = None + # Parameters the builder must pass through a callable before handing them over. + # batchgeneratorsv2 ranges are the reason: `BGContrast((0.7, 1.5))` samples 50/50 + # from [lo, 1] and [max(lo, 1), hi], where the bare tuple would be sampled + # uniformly -- a different distribution, not a formatting detail. The config + # stores the plain range; the adapter is applied on the way in. + param_adapters: Mapping[str, Callable[[Any], Any]] = field(default_factory=lambda: MappingProxyType({})) + # Constructor nudges that make the standalone smoke test exercise something. + smoke_kwargs: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({})) + summary: str = "" + + def __post_init__(self) -> None: + if self.cls.__name__ != self.name: + raise RegistryError( + f"registry name {self.name!r} does not match class name {self.cls.__name__!r}. " + "The config key is the class name, so these cannot diverge." + ) + if _has_var_keyword(self.cls) and self.forwards_to is None: + raise RegistryError( + f"{self.name}.__init__ takes **kwargs but declares no forwards_to. " + "Parameter validation reads the signature, so **kwargs would silently " + "accept anything. Name the parameters explicitly, or set forwards_to " + "to the class the kwargs are passed on to." + ) + + +def _constructor_signature(cls: type) -> inspect.Signature | None: + """The constructor signature, already without `self`. + + `inspect.signature(cls)` rather than `inspect.signature(cls.__init__)`: it drops + `self` for us, and mypy rejects reading `__init__` off a `type` as unsound. + """ + try: + return inspect.signature(cls) + except (TypeError, ValueError): # builtins and C extensions have no signature + return None + + +def _has_var_keyword(cls: type) -> bool: + """True if the constructor ends in **kwargs.""" + signature = _constructor_signature(cls) + if signature is None: + return False + return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in signature.parameters.values()) + + +def _named_params(cls: type) -> dict[str, inspect.Parameter]: + """Named keyword-assignable constructor parameters, minus *args/**kwargs.""" + signature = _constructor_signature(cls) + if signature is None: + return {} + skip = (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + return {name: p for name, p in signature.parameters.items() if p.kind not in skip} + + +# backend -> name -> entry. Insertion order is irrelevant; `entries()` sorts by PIPELINE_ORDER. +_REGISTRY: dict[Backend, dict[str, AugEntry]] = {backend: {} for backend in Backend} +# Held in a dict rather than a bare module global so `load_all` can flip it without +# a `global` statement. +_state: dict[str, bool] = {"loaded": False, "check_position": True} + +T = TypeVar("T", bound=type) + + +def register_entry(entry: AugEntry) -> AugEntry: + """Add an entry, rejecting anything that would make lookups ambiguous. + + The call form exists for third-party classes that cannot be decorated -- the + batchgeneratorsv2 transforms the CPU pipeline composes directly. + """ + backend_entries = _REGISTRY[entry.backend] + + clash = backend_entries.get(entry.name) + if clash is not None: + raise RegistryError(f"{entry.backend.value} augmentation {entry.name!r} is already registered (as {clash.cls.__module__}).") + + if _state["check_position"] and entry.name not in PIPELINE_ORDER[entry.backend]: + raise RegistryError( + f"{entry.backend.value} augmentation {entry.name!r} is not in PIPELINE_ORDER, so the pipeline " + f"has nowhere to put it. Add it to the tuple in registry.py, at the position it should run." + ) + + backend_entries[entry.name] = entry + return entry + + +def register( + *, + aug_id: AugId, + backend: Backend, + group: AugType, + forwards_to: type | None = None, + context_params: tuple[str, ...] = (), + wrap_random: bool = True, + force_sequential: bool = False, + external_asset: str | None = None, + smoke_kwargs: Mapping[str, Any] | None = None, + param_adapters: Mapping[str, Callable[[Any], Any]] | None = None, + summary: str = "", +) -> Callable[[T], T]: + """Register the decorated augmentation class. Returns the class unchanged.""" + + def decorate(cls: T) -> T: + register_entry( + AugEntry( + name=cls.__name__, + cls=cls, + backend=backend, + aug_id=aug_id, + group=group, + forwards_to=forwards_to, + context_params=context_params, + wrap_random=wrap_random, + force_sequential=force_sequential, + external_asset=external_asset, + smoke_kwargs=MappingProxyType(dict(smoke_kwargs or {})), + param_adapters=MappingProxyType(dict(param_adapters or {})), + summary=summary or _first_docstring_line(cls), + ) + ) + return cls + + return decorate + + +def _first_docstring_line(cls: type) -> str: + doc = inspect.getdoc(cls) or "" + return doc.strip().split("\n", 1)[0] + + +def load_all() -> None: + """Import every transform module so the decorators have run. Idempotent.""" + if _state["loaded"]: + return + # Set before importing, not after: a transform module that queries the registry at + # import time would otherwise re-enter this and recurse. + _state["loaded"] = True + importlib.import_module("smauglab.transforms") + + +def _ensure_loaded() -> None: + if not _state["loaded"]: + load_all() + + +def entries( + backend: Backend | None = None, + group: AugType | None = None, + aug_id: AugId | None = None, +) -> list[AugEntry]: + """Matching entries, in pipeline order. + + Order is a property of the augmentation, not of whoever last edited a config, + so it comes from the registry rather than from config key order. + """ + _ensure_loaded() + backends = [backend] if backend is not None else list(Backend) + found = [e for b in backends for e in _REGISTRY[b].values()] + if group is not None: + found = [e for e in found if e.group is group] + if aug_id is not None: + found = [e for e in found if e.aug_id is aug_id] + return sorted(found, key=lambda e: (e.backend.value, pipeline_position(e))) + + +def names(backend: Backend | None = None, group: AugType | None = None) -> list[str]: + """Registered config keys, in pipeline order.""" + return [e.name for e in entries(backend=backend, group=group)] + + +def get(name: str, backend: Backend | None = None) -> AugEntry: + """Resolve a config key to its entry, or raise with a suggestion.""" + _ensure_loaded() + backends = [backend] if backend is not None else list(Backend) + for b in backends: + entry = _REGISTRY[b].get(name) + if entry is not None: + return entry + raise UnknownAugmentationError(_unknown_augmentation_message(name, backend)) + + +def _unknown_augmentation_message(name: str, backend: Backend | None) -> str: + candidates = names(backend) + where = f"{backend.value} " if backend is not None else "" + lines = [f"unknown {where}augmentation {name!r}."] + + # Two kinds of near-miss, and they find different things. difflib catches + # typos; matching on the distinctive middle of the name catches a renamed key + # such as "ScharrTransform" -> "RandomScharrGPU", which shares too little with + # its replacement for any usable cutoff. Merged rather than used as a fallback: + # searching across backends, difflib alone would fill the list with CPU + # candidates and crowd out the GPU rename the caller is probably after. + stem = name.removeprefix("Random").removesuffix("Transform").removesuffix("GPU") + close = difflib.get_close_matches(name, candidates, n=3, cutoff=0.6) + close += [c for c in candidates if stem and stem.lower() in c.lower() and c not in close] + if close: + lines.append(f" Did you mean: {', '.join(close[:4])}?") + + lines.append(f" {len(candidates)} registered: run `smauglab list` to see them.") + return "\n".join(lines) + + +def accepted_params(entry: AugEntry) -> dict[str, inspect.Parameter]: + """Every parameter a config may set for this augmentation. + + The constructor signature is the whole truth -- there is no hand-maintained + schema to drift out of sync. `forwards_to` widens it where a class genuinely + passes kwargs on; `context_params` narrows it where the builder supplies the + value; and `p` is added for CPU entries whose probability lives on the + RandomTransform wrapper rather than on the transform itself. + """ + params = _named_params(entry.cls) + if entry.forwards_to is not None: + params = {**_named_params(entry.forwards_to), **params} + for name in entry.context_params: + params.pop(name, None) + if entry.backend is Backend.CPU and entry.wrap_random and "p" not in params: + params["p"] = inspect.Parameter("p", inspect.Parameter.KEYWORD_ONLY, default=1.0, annotation=float) + return params + + +def required_params(entry: AugEntry) -> set[str]: + """Parameters with no default, which a config must therefore supply.""" + return {name for name, p in accepted_params(entry).items() if p.default is inspect.Parameter.empty} + + +def unknown_parameter_message(entry: AugEntry, name: str) -> str: + """Explain an unaccepted parameter, with a suggestion where one exists.""" + allowed = sorted(accepted_params(entry)) + lines = [f"{entry.name}: unknown parameter {name!r}."] + close = difflib.get_close_matches(name, allowed, n=3, cutoff=0.6) + if close: + lines.append(f" Did you mean: {', '.join(close)}?") + elif name in RENAMED_HINTS: + lines.append(f" {name!r} -> {RENAMED_HINTS[name]}") + if name in entry.context_params: + lines[-1:] = [f" {name!r} is supplied by the trainer at runtime and must not appear in the config."] + lines.append(f" Accepted: {', '.join(allowed)}") + return "\n".join(lines) + + +def matrix() -> dict[AugId, dict[Backend, AugEntry | None]]: + """Concept -> backend -> implementing entry, or None where there is none. + + Every `AugId` gets a row even with no implementations anywhere, because an + empty cell is exactly the thing worth seeing. + """ + _ensure_loaded() + table: dict[AugId, dict[Backend, AugEntry | None]] = {aug_id: dict.fromkeys(Backend) for aug_id in AugId} + for entry in entries(): + table[entry.aug_id][entry.backend] = entry + return table + + +def render_matrix(fmt: str = "md") -> str: + """The CPU/GPU/MONAI coverage table.""" + table = matrix() + if fmt not in {"md", "table"}: + raise ValueError(f"unknown format {fmt!r}; expected 'md' or 'table'") + + def cell(entry: AugEntry | None) -> str: + if entry is None: + return "—" if fmt == "md" else "-" + return f"`{entry.name}`" if fmt == "md" else entry.name + + def group_of(row: dict[Backend, AugEntry | None]) -> str: + found = next((e for e in row.values() if e is not None), None) + return found.group.value if found else "" + + header = ["Augmentation", "Group", *(b.value for b in Backend)] + rows = [[aug_id.value, group_of(row), *(cell(row[b]) for b in Backend)] for aug_id, row in table.items()] + + if fmt == "md": + out = ["| " + " | ".join(header) + " |", "| " + " | ".join("---" for _ in header) + " |"] + out += ["| " + " | ".join(r) + " |" for r in rows] + return "\n".join(out) + + widths = [max(len(r[i]) for r in [header, *rows]) for i in range(len(header))] + return "\n".join(" ".join(c.ljust(w) for c, w in zip(row, widths)).rstrip() for row in [header, *rows]) + + +def _json_safe(value: Any) -> Any: + """Best-effort conversion of a default into something JSON can hold.""" + if value is inspect.Parameter.empty: + # Required: no default to show. Null makes the hole visible, and strict + # loading will reject it until someone fills it in. + return None + if isinstance(value, tuple): + return [_json_safe(v) for v in value] + if isinstance(value, list): + return [_json_safe(v) for v in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) + + +def render_template(backend: Backend) -> dict[str, dict[str, Any]]: + """A config section naming every registered augmentation at its defaults. + + Round-trips through the same `accepted_params`/`inspect.signature` path that + validation uses, so a template that fails to load is a real bug rather than a + documentation slip. Checked into the repo, which is what makes "an augmentation + exists but no config can reach it" a test failure. + """ + section: dict[str, dict[str, Any]] = {} + for entry in entries(backend): + section[entry.name] = {name: _json_safe(param.default) for name, param in sorted(accepted_params(entry).items())} + return section + + +def clear(mark_loaded: bool = True) -> None: + """Drop every entry. For tests only -- never call this from library code. + + `mark_loaded` defaults to True so a test that clears the registry and registers + its own synthetic entries does not have the real augmentations imported back in + underneath it on the next lookup. Pass False to restore normal lazy loading. + + Note this is NOT undoable on its own: `load_all()` re-imports + `smauglab.transforms`, but the module is already in `sys.modules` by then, so the + `@register` decorators do not run a second time and the registry would stay empty + for the rest of the process. Use `isolated()` unless you mean that. + """ + for backend_entries in _REGISTRY.values(): + backend_entries.clear() + _state["loaded"] = mark_loaded + + +@contextlib.contextmanager +def isolated(order: Mapping[Backend, tuple[str, ...]] | None = None) -> Iterator[None]: + """Empty the registry for the duration of the block, then put it back. + + For tests that register synthetic augmentations and must not see the real ones. + Restores by hand rather than by reloading, because re-importing the transform + modules would not re-run their decorators -- see `clear()`. + + `order` replaces PIPELINE_ORDER for the block, and registration is checked against + it. Omitting it disables the position check for the block instead, which is what + most tests want: they register throwaway classes whose pipeline position is not the + thing under test. `test_registering_a_class_absent_from_pipeline_order_fails` + passes an explicit table to cover the check itself. + """ + global PIPELINE_ORDER # noqa: PLW0603 -- swapping the table is the point + saved = {backend: dict(entries_) for backend, entries_ in _REGISTRY.items()} + saved_order = PIPELINE_ORDER + was_loaded = _state["loaded"] + clear(mark_loaded=True) + saved_check = _state["check_position"] + if order is None: + _state["check_position"] = False + else: + # Explicitly back on: these blocks nest inside an outer isolated() that turned + # the check off, and the point of passing a table is to be checked against it. + _state["check_position"] = True + PIPELINE_ORDER = MappingProxyType({backend: tuple(order.get(backend, ())) for backend in Backend}) + try: + yield + finally: + _state["check_position"] = saved_check + PIPELINE_ORDER = saved_order + for backend, entries_ in saved.items(): + _REGISTRY[backend].clear() + _REGISTRY[backend].update(entries_) + _state["loaded"] = was_loaded diff --git a/smauglab/transforms/__init__.py b/smauglab/transforms/__init__.py new file mode 100644 index 0000000..91a5776 --- /dev/null +++ b/smauglab/transforms/__init__.py @@ -0,0 +1,9 @@ +"""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. + +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. +""" diff --git a/smauglab/transforms/cpu/__init__.py b/smauglab/transforms/cpu/__init__.py new file mode 100644 index 0000000..0ca92e8 --- /dev/null +++ b/smauglab/transforms/cpu/__init__.py @@ -0,0 +1 @@ +"""CPU augmentations, built on batchgeneratorsv2 transforms.""" diff --git a/smauglab/transforms/gpu/__init__.py b/smauglab/transforms/gpu/__init__.py new file mode 100644 index 0000000..187f729 --- /dev/null +++ b/smauglab/transforms/gpu/__init__.py @@ -0,0 +1 @@ +"""GPU augmentations, built on kornia's 3D augmentation base classes.""" diff --git a/smauglab/utils/__init__.py b/smauglab/utils/__init__.py new file mode 100644 index 0000000..0c90b86 --- /dev/null +++ b/smauglab/utils/__init__.py @@ -0,0 +1,13 @@ +"""NIfTI image handling. + +Only `image.py` lives here. `utils.py` used to sit alongside it -- MONAI training-loop +helpers, argparse tuple parsers, a Dice function -- but nothing under `smauglab/` +imported any of it once the `__main__` demo blocks moved out to `scripts/`, so it +shipped in every wheel for the benefit of two standalone scripts. It is now +`scripts/_common.py`. + +`image.py` stayed despite having no in-package consumer either: five modules in the +sibling segtransferaug repository import `smauglab.utils.image.Image`, so it is part +of the public API in practice. It is a vendored subset of spinalcordtoolbox's +`image.py` -- see the class docstrings for the upstream links. +""" diff --git a/unit_tests/test_registry.py b/unit_tests/test_registry.py new file mode 100644 index 0000000..b4c78e8 --- /dev/null +++ b/unit_tests/test_registry.py @@ -0,0 +1,260 @@ +"""The registry machinery, exercised against synthetic transform classes. + +Deliberately not against the real augmentations: those get registered in a later +stage, and these tests are about the mechanism -- registration invariants, +signature-derived parameter validation, and the coverage matrix. Keeping them +synthetic means they stay fast and cannot break for reasons unrelated to the +registry itself. + +`registry.clear()` empties the global table, so every test here builds the world +it needs and tears it down again. +""" + +from __future__ import annotations + +import inspect +import unittest + +from smauglab import registry +from smauglab.registry import ( + AugEntry, + AugId, + AugType, + Backend, + RegistryError, + UnknownAugmentationError, +) + + +class RegistryTestCase(unittest.TestCase): + """Isolates each test from the global registry.""" + + def setUp(self) -> None: + # Empties the registry so these tests see only what they register, and puts + # the real augmentations back afterwards. A bare clear() cannot be undone: + # load_all() re-imports an already-imported module, so the decorators never + # run again and every later test in the process would see an empty registry. + context = registry.isolated() + context.__enter__() + self.addCleanup(context.__exit__, None, None, None) + + +def make_transform(name: str, **params): + """A throwaway class whose __init__ has exactly the given parameters.""" + defaults = ", ".join(f"{k}={v!r}" for k, v in params.items()) + namespace: dict = {} + exec( # noqa: S102 -- building a signature is the point of this helper + f"def __init__(self, {defaults}): pass" if defaults else "def __init__(self): pass", + namespace, + ) + return type(name, (), {"__init__": namespace["__init__"], "__doc__": f"{name} summary line.\n\nMore."}) + + +class TestRegistration(RegistryTestCase): + def test_decorator_registers_and_returns_the_class(self): + cls = make_transform("RandomThingGPU", p=1.0) + decorated = registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(cls) + + self.assertIs(decorated, cls, "the decorator must not replace the class") + self.assertEqual(registry.get("RandomThingGPU", Backend.GPU).cls, cls) + + def test_summary_defaults_to_the_first_docstring_line(self): + cls = make_transform("RandomThingGPU", p=1.0) + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(cls) + self.assertEqual(registry.get("RandomThingGPU").summary, "RandomThingGPU summary line.") + + def test_name_must_match_the_class_name(self): + """The config key IS the class name, so a mismatch has to be impossible.""" + with self.assertRaises(RegistryError) as caught: + AugEntry( + name="SomethingElse", + cls=make_transform("RandomThingGPU"), + backend=Backend.GPU, + aug_id=AugId.SCHARR, + group=AugType.TA, + ) + self.assertIn("does not match class name", str(caught.exception)) + + def test_duplicate_name_is_rejected(self): + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomThingGPU", p=1.0)) + with self.assertRaises(RegistryError) as caught: + registry.register(aug_id=AugId.LAPLACE, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomThingGPU", p=1.0)) + self.assertIn("already registered", str(caught.exception)) + + def test_registering_a_class_absent_from_pipeline_order_fails(self): + """PIPELINE_ORDER is where a transform's pipeline position lives, so a class + missing from it has nowhere to run. That is an error at import, not a silent + append, which is what keeps the table from falling out of date.""" + order = {Backend.GPU: ("RandomListedGPU",)} + with registry.isolated(order=order): + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomListedGPU", p=1.0)) + with self.assertRaises(registry.RegistryError) as caught: + registry.register(aug_id=AugId.LAPLACE, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomUnlistedGPU", p=1.0)) + self.assertIn("PIPELINE_ORDER", str(caught.exception)) + + def test_pipeline_order_positions_follow_the_table(self): + order = {Backend.GPU: ("RandomSecondGPU", "RandomFirstGPU")} + with registry.isolated(order=order): + # Registered in the opposite order to the table, to prove the table wins. + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomFirstGPU", p=1.0)) + registry.register(aug_id=AugId.LAPLACE, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomSecondGPU", p=1.0)) + self.assertEqual(registry.names(Backend.GPU), ["RandomSecondGPU", "RandomFirstGPU"]) + + def test_the_same_concept_can_exist_on_two_backends(self): + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomScharrGPU", p=1.0)) + registry.register(aug_id=AugId.SCHARR, backend=Backend.CPU, group=AugType.TA)(make_transform("ScharrConvTransform")) + self.assertEqual(len(registry.entries()), 2) + + def test_var_keyword_without_forwards_to_is_rejected(self): + """**kwargs would make signature-based validation accept anything.""" + + class RandomSloppyGPU: + def __init__(self, p: float = 1.0, **kwargs): + pass + + with self.assertRaises(RegistryError) as caught: + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(RandomSloppyGPU) + self.assertIn("**kwargs", str(caught.exception)) + + def test_var_keyword_is_allowed_when_forwards_to_is_declared(self): + target = make_transform("Generator", n_labels=3, blur=0.5) + + class RandomForwardingGPU: + def __init__(self, p: float = 1.0, **kwargs): + pass + + registry.register(aug_id=AugId.SYNTHSEG, backend=Backend.GPU, group=AugType.TA, forwards_to=target)(RandomForwardingGPU) + self.assertIn("n_labels", registry.accepted_params(registry.get("RandomForwardingGPU"))) + + +class TestLookup(RegistryTestCase): + def setUp(self) -> None: + super().setUp() + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)( + make_transform("RandomScharrGPU", p=1.0, absolute=True) + ) + registry.register(aug_id=AugId.FLIP, backend=Backend.GPU, group=AugType.GEO)(make_transform("RandomFlipTransformGPU", p=1.0)) + + def test_entries_come_back_in_pipeline_order_not_registration_order(self): + self.assertEqual(registry.names(Backend.GPU), ["RandomFlipTransformGPU", "RandomScharrGPU"]) + + def test_filtering_by_group(self): + self.assertEqual(registry.names(Backend.GPU, group=AugType.GEO), ["RandomFlipTransformGPU"]) + + def test_unknown_name_suggests_a_close_match(self): + with self.assertRaises(UnknownAugmentationError) as caught: + registry.get("RandomScharGPU", Backend.GPU) + self.assertIn("Did you mean: RandomScharrGPU?", str(caught.exception)) + + def test_legacy_key_fails_but_points_at_its_replacement(self): + """The hard break: old spellings must raise, not quietly work. + + difflib alone cannot bridge 'ScharrTransform' -> 'RandomScharrGPU' + (too little shared prefix), so the stem fallback has to carry it. + """ + with self.assertRaises(UnknownAugmentationError) as caught: + registry.get("ScharrTransform", Backend.GPU) + self.assertIn("Did you mean: RandomScharrGPU?", str(caught.exception)) + + def test_renamed_hints_are_diagnostics_only(self): + """Nothing in RENAMED_HINTS may be a working config key.""" + for stale in registry.RENAMED_HINTS: + with self.subTest(key=stale), self.assertRaises(UnknownAugmentationError): + registry.get(stale) + + +class TestAcceptedParams(RegistryTestCase): + def test_params_come_from_the_constructor_signature(self): + registry.register(aug_id=AugId.GAUSSIAN_NOISE, backend=Backend.GPU, group=AugType.GE)( + make_transform("RandomGaussianNoiseGPU", p=1.0, mean=0.0, std=0.1) + ) + entry = registry.get("RandomGaussianNoiseGPU") + self.assertEqual(set(registry.accepted_params(entry)), {"p", "mean", "std"}) + + def test_context_params_are_excluded(self): + registry.register( + aug_id=AugId.SPATIAL, + backend=Backend.CPU, + group=AugType.GEO, + wrap_random=False, + context_params=("rotation",), + )(make_transform("SpatialTransform", rotation=None, p_rotation=0.2)) + entry = registry.get("SpatialTransform") + self.assertNotIn("rotation", registry.accepted_params(entry)) + self.assertIn("p_rotation", registry.accepted_params(entry)) + + def test_cpu_wrapped_transforms_accept_p_even_though_the_class_does_not(self): + """batchgeneratorsv2 puts the probability on the RandomTransform wrapper.""" + registry.register(aug_id=AugId.SCHARR, backend=Backend.CPU, group=AugType.TA)(make_transform("ScharrConvTransform", absolute=True)) + entry = registry.get("ScharrConvTransform") + self.assertIn("p", registry.accepted_params(entry)) + + def test_unwrapped_cpu_transforms_do_not_get_a_synthetic_p(self): + registry.register(aug_id=AugId.MIRROR, backend=Backend.CPU, group=AugType.GEO, wrap_random=False)( + make_transform("MirrorTransform", allowed_axes=None) + ) + self.assertNotIn("p", registry.accepted_params(registry.get("MirrorTransform"))) + + def test_required_params_are_those_without_a_default(self): + """RandomAffineGPU really does take a required `degrees` today.""" + + class RandomAffineGPU: + def __init__(self, degrees, p: float = 1.0): + pass + + registry.register(aug_id=AugId.AFFINE, backend=Backend.GPU, group=AugType.GEO)(RandomAffineGPU) + self.assertEqual(registry.required_params(registry.get("RandomAffineGPU")), {"degrees"}) + + def test_unknown_parameter_message_suggests_p_for_probability(self): + registry.register(aug_id=AugId.GAUSSIAN_NOISE, backend=Backend.GPU, group=AugType.GE)( + make_transform("RandomGaussianNoiseGPU", p=1.0, std=0.1) + ) + message = registry.unknown_parameter_message(registry.get("RandomGaussianNoiseGPU"), "probability") + self.assertIn("'probability' -> p", message) + self.assertIn("Accepted: p, std", message) + + def test_unknown_parameter_message_flags_context_params_specifically(self): + registry.register( + aug_id=AugId.SPATIAL, + backend=Backend.CPU, + group=AugType.GEO, + wrap_random=False, + context_params=("rotation",), + )(make_transform("SpatialTransform", rotation=None)) + message = registry.unknown_parameter_message(registry.get("SpatialTransform"), "rotation") + self.assertIn("supplied by the trainer", message) + + +class TestMatrix(RegistryTestCase): + def test_every_aug_id_gets_a_row_even_with_no_implementations(self): + self.assertEqual(set(registry.matrix()), set(AugId)) + + def test_a_concept_joins_its_backends_into_one_row(self): + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomScharrGPU", p=1.0)) + registry.register(aug_id=AugId.SCHARR, backend=Backend.CPU, group=AugType.TA)(make_transform("ScharrConvTransform")) + row = registry.matrix()[AugId.SCHARR] + self.assertEqual(row[Backend.GPU].name, "RandomScharrGPU") + self.assertEqual(row[Backend.CPU].name, "ScharrConvTransform") + self.assertIsNone(row[Backend.MONAI], "no MONAI implementations exist yet") + + def test_markdown_render_marks_the_gap(self): + registry.register(aug_id=AugId.SCHARR, backend=Backend.GPU, group=AugType.TA)(make_transform("RandomScharrGPU", p=1.0)) + rendered = registry.render_matrix("md") + self.assertIn("| scharr | TA | `RandomScharrGPU` | — | — |", rendered) + + def test_unknown_format_is_rejected(self): + with self.assertRaises(ValueError): + registry.render_matrix("yaml") + + +class TestModuleHygiene(unittest.TestCase): + def test_registry_does_not_import_torch(self): + """`smauglab list` must be able to answer without paying for torch.""" + source = inspect.getsource(registry) + for banned in ("import torch", "import kornia", "from smauglab.transforms"): + with self.subTest(banned=banned): + self.assertNotIn(f"\n{banned}", source, f"registry.py must not import {banned!r} at module scope") + + +if __name__ == "__main__": + unittest.main()