claude refactor of palette to enable exchange of building blocks, req…

auglab/transforms/gpu/palette/__init__.py

@@ -0,0 +1,17 @@
+from auglab.transforms.gpu.palette.base import (
+    BlockContext,
+    InitialPartitioner,
+    RefinementPartitioner,
+    signed_alpha_affine_remap,
+)
+from auglab.transforms.gpu.palette.factory import build_palette_from_cfg
+from auglab.transforms.gpu.palette.transform import PaletteSynthesisGPU
+
+__all__ = [
+    "BlockContext",
+    "InitialPartitioner",
+    "RefinementPartitioner",
+    "signed_alpha_affine_remap",
+    "PaletteSynthesisGPU",
+    "build_palette_from_cfg",
+]

auglab/transforms/gpu/palette/base.py

@@ -0,0 +1,65 @@
+"""Base contracts and shared helpers for the composable PALETTE pipeline."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+from torch import nn
+
+
+@dataclass
+class BlockContext:
+    """Per-sample state shared across partitioners and remap steps."""
+    image01: torch.Tensor          # (N,) float, min-max normalized to [0,1]
+    fg_mask: torch.Tensor          # (N,) float 0/1
+    coords: torch.Tensor           # (N, 3) ijk voxel coords
+    shape: Tuple[int, int, int]    # (D, H, W)
+    device: torch.device
+
+
+class InitialPartitioner(nn.Module):
+    """Produce the initial region map from the raw image."""
+
+    def partition(self, ctx: BlockContext) -> Tuple[torch.Tensor, int]:
+        raise NotImplementedError
+
+
+class RefinementPartitioner(nn.Module):
+    """Subdivide an existing region map. May optionally consult the image."""
+
+    def refine(
+        self,
+        ctx: BlockContext,
+        region_ids: torch.Tensor,
+        n_regions: int,
+    ) -> Tuple[torch.Tensor, int]:
+        raise NotImplementedError
+
+
+def signed_alpha_affine_remap(
+    image01: torch.Tensor,
+    fg_mask: torch.Tensor,
+    region_ids: torch.Tensor,
+    n_regions: int,
+    alpha_magnitude_range: Tuple[float, float],
+    eps: float = 1e-7,
+) -> torch.Tensor:
+    """Signed-alpha per-region affine remap: y = μ_c + α_c · (x − mean_c).
+
+    Mirrors fromSeg.py:392-401. Foreground-only means; output clamped to [0,1]
+    and multiplied by the foreground mask.
+    """
+    device = image01.device
+    alpha_lo, alpha_hi = alpha_magnitude_range
+
+    s_c = torch.zeros(n_regions, device=device).scatter_add_(0, region_ids, image01 * fg_mask)
+    n_c = torch.zeros(n_regions, device=device).scatter_add_(0, region_ids, fg_mask)
+    mean_c = s_c / n_c.clamp(min=eps)
+
+    mu_c = torch.rand(n_regions, device=device)
+    mag_c = torch.rand(n_regions, device=device) * (alpha_hi - alpha_lo) + alpha_lo
+    sign_c = (torch.rand(n_regions, device=device) > 0.5).float() * 2 - 1
+    alp_c = mag_c * sign_c
+
+    return (mu_c[region_ids] + alp_c[region_ids] * (image01 - mean_c[region_ids])).clamp(0, 1) * fg_mask

auglab/transforms/gpu/palette/factory.py

@@ -0,0 +1,113 @@
+"""Config → PaletteSynthesisGPU factory with typed partitioner registries."""
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from auglab.transforms.gpu.palette.overlay import AnatomicalLabelOverlay
+from auglab.transforms.gpu.palette.partitioners import (
+    EMGMMInitial,
+    EMGMMRefinement,
+    IdentityRefinement,
+    KMeans1DInitial,
+    VoronoiRefinement,
+)
+from auglab.transforms.gpu.palette.transform import PaletteSynthesisGPU
+
+
+INITIAL_REGISTRY: Dict[str, type] = {
+    "kmeans1d": KMeans1DInitial,
+    "em_gmm":   EMGMMInitial,
+}
+
+REFINEMENT_REGISTRY: Dict[str, type] = {
+    "voronoi":  VoronoiRefinement,
+    "em_gmm":   EMGMMRefinement,
+    "identity": IdentityRefinement,
+}
+
+# Top-level config keys consumed directly by PaletteSynthesisGPU (everything
+# else must be routed into a nested block or is rejected as unknown).
+_TOP_KEYS = {
+    "alpha_magnitude_range",
+    "dark_threshold",
+    "blur_sigmas_pre",
+    "blur_sigmas_post",
+    "p",
+}
+
+# Keys accepted at the top of the config but not passed to the transform ctor.
+_TOP_META_KEYS = {"probability", "initial_partitioner", "refinement_partitioners", "overlay"}
+
+
+def build_palette_from_cfg(cfg: Dict[str, Any]) -> PaletteSynthesisGPU:
+    """Build a ``PaletteSynthesisGPU`` from a nested config dict.
+
+    Raises ``ValueError`` if the initial partitioner is missing, if a type
+    string is placed in the wrong registry slot, or if unknown top-level keys
+    are present.
+    """
+    if "initial_partitioner" not in cfg:
+        raise ValueError(
+            "PaletteSynthesisTransform config requires an 'initial_partitioner' block "
+            f"(one of: {list(INITIAL_REGISTRY)})"
+        )
+
+    unknown = set(cfg) - _TOP_KEYS - _TOP_META_KEYS
+    if unknown:
+        raise ValueError(
+            f"Unknown top-level keys in PaletteSynthesisTransform: {sorted(unknown)}. "
+            f"Expected any of: {sorted(_TOP_KEYS | _TOP_META_KEYS)}"
+        )
+
+    init_cfg = dict(cfg["initial_partitioner"])
+    init_type = init_cfg.pop("type", None)
+    if init_type is None:
+        raise ValueError("initial_partitioner must specify a 'type' field")
+    if init_type in REFINEMENT_REGISTRY and init_type not in INITIAL_REGISTRY:
+        raise ValueError(
+            f"'{init_type}' is a refinement partitioner, not an initial one. "
+            f"Move it into 'refinement_partitioners'. "
+            f"Valid initial types: {list(INITIAL_REGISTRY)}"
+        )
+    if init_type not in INITIAL_REGISTRY:
+        raise ValueError(
+            f"Unknown initial partitioner type '{init_type}'. "
+            f"Valid types: {list(INITIAL_REGISTRY)}"
+        )
+    initial = INITIAL_REGISTRY[init_type](**init_cfg)
+
+    refinements = []
+    for i, r_cfg in enumerate(cfg.get("refinement_partitioners", []) or []):
+        r_cfg = dict(r_cfg)
+        r_type = r_cfg.pop("type", None)
+        if r_type is None:
+            raise ValueError(f"refinement_partitioners[{i}] must specify a 'type' field")
+        if r_type in INITIAL_REGISTRY and r_type not in REFINEMENT_REGISTRY:
+            raise ValueError(
+                f"'{r_type}' is an initial partitioner and cannot be used as a refinement. "
+                f"Valid refinement types: {list(REFINEMENT_REGISTRY)}"
+            )
+        if r_type not in REFINEMENT_REGISTRY:
+            raise ValueError(
+                f"Unknown refinement partitioner type '{r_type}'. "
+                f"Valid types: {list(REFINEMENT_REGISTRY)}"
+            )
+        refinements.append(REFINEMENT_REGISTRY[r_type](**r_cfg))
+
+    ov_cfg = cfg.get("overlay")
+    overlay = None
+    if ov_cfg is not None:
+        ov_kwargs = {k: v for k, v in ov_cfg.items() if k != "enabled"}
+        if ov_cfg.get("enabled", True):
+            overlay = AnatomicalLabelOverlay(**ov_kwargs)
+
+    top_kwargs = {k: cfg[k] for k in _TOP_KEYS if k in cfg}
+    if "probability" in cfg and "p" not in top_kwargs:
+        top_kwargs["p"] = cfg["probability"]
+
+    return PaletteSynthesisGPU(
+        initial_partitioner=initial,
+        refinement_partitioners=refinements,
+        overlay=overlay,
+        **top_kwargs,
+    )

auglab/transforms/gpu/palette/overlay.py

@@ -0,0 +1,108 @@
+"""Per-anatomical-label overlay: fixed algorithm, tunable frequency and blend."""
+from __future__ import annotations
+
+from typing import List, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+BlendSpec = Union[float, Sequence[float]]
+
+
+class AnatomicalLabelOverlay(nn.Module):
+    """PALETTE per-anatomical-label affine remap with a tunable blend.
+
+    Reproduces fromSeg.py:414-447 algorithmically. For each foreground label,
+    with probability ``label_remap_prob``, samples a fresh (μ, α) and computes
+    ``new_vals = μ + α·(synth − mean_label)``, then blends into the current
+    synthesised image via a per-label blend strength.
+
+    Args:
+        label_remap_prob: per-label per-sample probability of applying the remap.
+        min_label_voxels: minimum voxel count for a label to be eligible.
+        label_classes: if set, restrict overlay to these class indices.
+        blend_strength: scalar in [0,1] (full overwrite = 1.0) or [lo, hi]
+            sampled per label per sample. Applied as
+            ``write_mask = c_mask · apply · blend``.
+        alpha_magnitude_range: [lo, hi] for |α| of the signed-alpha remap.
+    """
+
+    def __init__(
+        self,
+        label_remap_prob: float = 0.5,
+        min_label_voxels: int = 4,
+        label_classes: Optional[List[int]] = None,
+        blend_strength: BlendSpec = 1.0,
+        alpha_magnitude_range: Sequence[float] = (0.5, 2.0),
+    ) -> None:
+        super().__init__()
+        self.label_remap_prob = float(label_remap_prob)
+        self.min_label_voxels = int(min_label_voxels)
+        self.label_classes = None if label_classes is None else list(label_classes)
+        self.blend_strength = _normalize_blend(blend_strength)
+        self.alpha_magnitude_range = tuple(alpha_magnitude_range)
+
+    def apply(
+        self,
+        synth: torch.Tensor,        # (B, N)
+        labels: torch.Tensor,       # (B, 1, D, H, W) long
+        shape: Tuple[int, int, int],
+    ) -> torch.Tensor:
+        B, N = synth.shape
+        device = synth.device
+        alpha_lo, alpha_hi = self.alpha_magnitude_range
+        blend_lo, blend_hi = self.blend_strength
+        D, H, W = shape
+
+        if labels.shape[2:] != (D, H, W):
+            labels = F.interpolate(labels.float(), size=(D, H, W), mode="nearest").long()
+        lbl = labels[:, 0].reshape(B, N).clamp(min=0)
+
+        unique_classes = lbl.unique()
+        unique_classes = unique_classes[unique_classes > 0]
+        if self.label_classes is not None:
+            keep = torch.tensor(self.label_classes, device=device)
+            unique_classes = unique_classes[torch.isin(unique_classes, keep)]
+
+        for c in unique_classes:
+            c_val = int(c.item())
+            c_mask = (lbl == c_val).float()
+            c_cnt = c_mask.sum(dim=1, keepdim=True)
+
+            apply = (
+                (torch.rand(B, 1, device=device) < self.label_remap_prob)
+                & (c_cnt >= self.min_label_voxels)
+            ).float()
+
+            if apply.sum() == 0:
+                continue
+
+            c_mean = (synth * c_mask).sum(dim=1, keepdim=True) / c_cnt.clamp(min=1)
+
+            mu_c = torch.rand(B, 1, device=device)
+            mag_c = torch.rand(B, 1, device=device) * (alpha_hi - alpha_lo) + alpha_lo
+            sign_c = (torch.rand(B, 1, device=device) > 0.5).float() * 2 - 1
+            alp_c = mag_c * sign_c
+
+            if blend_lo == blend_hi:
+                blend = torch.full((B, 1), blend_lo, device=device)
+            else:
+                blend = torch.rand(B, 1, device=device) * (blend_hi - blend_lo) + blend_lo
+
+            new_vals = (mu_c + alp_c * (synth - c_mean)).clamp(0, 1)
+            write_mask = c_mask * apply * blend
+            synth = synth * (1.0 - write_mask) + new_vals * write_mask
+
+        return synth
+
+
+def _normalize_blend(spec: BlendSpec) -> Tuple[float, float]:
+    if isinstance(spec, (int, float)):
+        v = float(spec)
+        return (v, v)
+    lo, hi = float(spec[0]), float(spec[1])
+    if lo > hi:
+        raise ValueError(f"blend_strength range must be non-decreasing, got [{lo}, {hi}]")
+    return (lo, hi)