Transforms

class Transform(nn.Module):
    """Abstract class for all TorchIO transforms.

    When called, the input can be an instance of
    [`Subject`][torchio.Subject],
    [`Image`][torchio.Image],
    [`torch.Tensor`][torch.Tensor],
    [`numpy.ndarray`][numpy.ndarray],
    [`SimpleITK.Image`](https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1Image.html),
    [`nibabel.Nifti1Image`](https://nipy.org/nibabel/reference/nibabel.nifti1.html),
    [`dict`][dict] containing 4D tensors as values,
    [`ImagesBatch`][torchio.ImagesBatch], or
    [`SubjectsBatch`][torchio.SubjectsBatch].
    The output type always matches the input type.

    All subclasses must override
    [`apply_transform()`][torchio.Transform.apply_transform],
    which receives a [`SubjectsBatch`][torchio.SubjectsBatch] and
    returns the transformed batch.

    Args:
        p: Probability that this transform will be applied.
        copy: Make a deep copy of the input before applying the
            transform. When transforms are composed with
            [`Compose`][torchio.Compose], the outer `Compose`
            copies once and sets `copy=False` on inner transforms
            to avoid redundant copies.
        include: Sequence of strings with the names of the only images
            to which the transform will be applied.
        exclude: Sequence of strings with the names of the images to
            which the transform will *not* be applied.
    """

    def __init__(
        self,
        *,
        p: float = 1.0,
        copy: bool = True,
        include: list[str] | None = None,
        exclude: list[str] | None = None,
    ) -> None:
        super().__init__()
        if not 0 <= p <= 1:
            msg = f"Probability must be in [0, 1], got {p}"
            raise ValueError(msg)
        self.p = p
        self.copy = copy
        self.include = include
        self.exclude = exclude

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        _TRANSFORM_REGISTRY[cls.__name__] = cls

    def _warn_if_noop(self, *, is_noop: bool, hint: str) -> None:
        """Warn that the transform leaves the data unchanged.

        Augmentation transforms whose parameters are sampled from a
        range default to an identity (no-op) when constructed with no
        arguments, so that randomness must be requested explicitly. This
        warns the user when that happens (or whenever the given
        parameters produce a no-op).

        Args:
            is_noop: Whether the configured transform is an identity.
            hint: Example argument to suggest in the warning message.
        """
        if is_noop:
            warnings.warn(
                f"{type(self).__name__} is a no-op with the given parameters"
                " and will not change the data. Pass arguments to apply an"
                f" effect (e.g. {hint}), or a range like (a, b) for random"
                " augmentation.",
                stacklevel=3,
            )

    def __repr__(self) -> str:
        """Show only non-default fields for a compact repr."""
        from .parameter_range import _ParameterRange

        parts = []
        for name, default in _collect_init_params(type(self)).items():
            value = getattr(self, name, default)
            if isinstance(value, _ParameterRange):
                if value._original == default:
                    continue
            elif value == default:
                continue
            parts.append(f"{name}={value!r}")
        return f"{type(self).__name__}({', '.join(parts)})"

    def __add__(self, other: object) -> Transform:
        """Compose two transforms: `t1 + t2` → `Compose([t1, t2])`."""
        if not isinstance(other, Transform):
            return NotImplemented
        from .compose import Compose

        left = self.transforms if isinstance(self, Compose) else [self]
        right = other.transforms if isinstance(other, Compose) else [other]
        return Compose([*left, *right])

    def __or__(self, other: object) -> Transform:
        """Random choice: `t1 | t2` → `OneOf([t1, t2])`."""
        if not isinstance(other, Transform):
            return NotImplemented
        from .compose import OneOf

        left = self.transforms if isinstance(self, OneOf) else [self]
        right = other.transforms if isinstance(other, OneOf) else [other]
        return OneOf([*left, *right])

    @overload
    def forward(self, data: Subject) -> Subject: ...
    @overload
    def forward(self, data: Image) -> Image: ...
    @overload
    def forward(self, data: Tensor) -> Tensor: ...
    @overload
    def forward(self, data: np.ndarray) -> np.ndarray: ...
    @overload
    def forward(self, data: sitk.Image) -> sitk.Image: ...
    @overload
    def forward(self, data: nib.Nifti1Image) -> nib.Nifti1Image: ...
    @overload
    def forward(self, data: dict) -> dict: ...
    @overload
    def forward(self, data: ImagesBatch) -> ImagesBatch: ...
    @overload
    def forward(self, data: SubjectsBatch) -> SubjectsBatch: ...

    def forward(self, data: Any) -> Any:
        """Apply the transform.

        The output type always matches the input type.

        Args:
            data: Input data to transform.
        """
        if self.copy:
            data = _copy.deepcopy(data)
        batch, unwrap = self._wrap(data)
        if torch.rand(1).item() > self.p:
            return unwrap(batch)
        params = self.make_params(batch)
        batch = self.apply_transform(batch, params)
        # Record history on the batch
        trace = AppliedTransform(name=type(self).__name__, params=params)
        if not hasattr(batch, "applied_transforms"):
            batch.applied_transforms = []
        batch.applied_transforms.append(trace)
        result = unwrap(batch)
        # Propagate history to outputs that can carry it
        if (
            hasattr(batch, "applied_transforms")
            and not isinstance(result, (SubjectsBatch, Tensor, np.ndarray))
            and not isinstance(result, dict)
        ):
            with contextlib.suppress(AttributeError):
                result.applied_transforms = list(batch.applied_transforms)
        return result

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        """Sample random parameters for this transform.

        Override in subclasses that have random behavior.

        Args:
            batch: A `SubjectsBatch`.

        Returns:
            Dict of sampled parameters.
        """
        return {}

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        """Apply the transform with the given parameters.

        Must be overridden by subclasses. Receives a `SubjectsBatch`
        whose `ImagesBatch` entries contain 5D tensors
        `(B, C, I, J, K)`. Use negative indexing (`-3`, `-2`,
        `-1`) for spatial dims.

        Args:
            batch: A `SubjectsBatch` to transform.
            params: Parameters from `make_params`.

        Returns:
            Transformed `SubjectsBatch`.
        """
        raise NotImplementedError

    @property
    def invertible(self) -> bool:
        """Whether this transform can be inverted."""
        return False

    def inverse(self, params: dict[str, Any]) -> Transform:
        """Return a transform that undoes this one.

        Override in invertible subclasses. The returned transform,
        when applied, reverses the effect of the forward pass with
        the given parameters.

        Args:
            params: The parameters recorded in the forward pass.

        Returns:
            A new `Transform` instance that inverts this one.
        """
        msg = f"{type(self).__name__} is not invertible"
        raise NotImplementedError(msg)

    def _get_images(self, batch: SubjectsBatch) -> dict[str, ImagesBatch]:
        """Get image batches filtered by include/exclude."""
        images = batch.images
        if self.include is not None:
            images = {k: v for k, v in images.items() if k in self.include}
        if self.exclude is not None:
            images = {k: v for k, v in images.items() if k not in self.exclude}
        return images

    def to_hydra(self) -> dict[str, Any]:
        """Export as a Hydra-compatible config dict.

        Returns a dict with `_target_` set to the fully qualified
        class name and only non-default field values included.

        Returns:
            Dict suitable for `hydra.utils.instantiate()`.
        """
        from .parameter_range import _ParameterRange

        cls = type(self)
        target = f"torchio.{cls.__qualname__}"
        cfg: dict[str, Any] = {"_target_": target}

        for name, default in _collect_init_params(cls).items():
            value = getattr(self, name, default)
            if isinstance(value, _ParameterRange):
                if value._original == default:
                    continue
                value = _hydra_value(value._original)
            elif value == default:
                continue
            else:
                value = _hydra_value(value)
            cfg[name] = value
        return cfg

    @staticmethod
    def _wrap(
        data: Any,
    ) -> tuple[Any, Any]:
        """Wrap any input into a SubjectsBatch; return (batch, unwrap_fn)."""
        from ..data.batch import ImagesBatch
        from ..data.batch import SubjectsBatch

        match data:
            case SubjectsBatch():
                return data, _unwrap_subjects_batch
            case ImagesBatch():
                sb = SubjectsBatch({"tio_default_image": data})
                return sb, _unwrap_images_batch
            case Subject():
                sb = SubjectsBatch.from_subjects([data])
                return sb, _unwrap_subject
            case dict():
                return _wrap_dict(data)
            case _:
                return _wrap_scalar_input(data)

class SpatialTransform(Transform):
    """Base for transforms that modify spatial geometry.

    Spatial transforms apply to all images (ScalarImage and LabelMap),
    and also transform any Points and BoundingBoxes attached to the
    Subject.
    """

class IntensityTransform(Transform):
    """Base for transforms that modify voxel intensities.

    Intensity transforms apply only to `ScalarImage` instances,
    leaving `LabelMap` and annotations unchanged.
    """

    def _get_images(self, batch: SubjectsBatch) -> dict[str, ImagesBatch]:
        """Filter to ScalarImage batches only, then apply include/exclude."""
        images = {
            k: v for k, v in batch.images.items() if v._image_class is ScalarImage
        }
        if self.include is not None:
            images = {k: v for k, v in images.items() if k in self.include}
        if self.exclude is not None:
            images = {k: v for k, v in images.items() if k not in self.exclude}
        return images

@dataclass
class AppliedTransform:
    """Record of a transform application, stored in Subject history.

    Attributes:
        name: Class name of the transform.
        params: Sampled parameters (JSON-serializable).
    """

    name: str
    params: dict[str, Any] = field(default_factory=dict)

class To(Transform):
    """Move all data to a device and/or cast to a dtype.

    Wraps the `to()` method as a transform so it can be used inside
    [`Compose`][torchio.Compose] pipelines.

    Args:
        *to_args: Positional arguments forwarded to
            [`torch.Tensor.to()`](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html).
            Typically a device string (`"cpu"`, `"cuda"`,
            `"mps"`) or a `torch.dtype` (`torch.float16`).
        **to_kwargs: Keyword arguments forwarded to
            `torch.Tensor.to()`.

    Examples:
        >>> import torchio as tio
        >>> transform = tio.To(torch.float16)
        >>> transform = tio.To("cuda")
        >>> pipeline = tio.Compose([
        ...     tio.To("cuda"),
        ...     tio.Noise(std=0.1),
        ... ])
    """

    def __init__(self, *to_args: Any, **to_kwargs: Any) -> None:
        super().__init__()
        self.to_args = to_args
        self.to_kwargs = to_kwargs

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        return {"to_args": self.to_args, "to_kwargs": self.to_kwargs}

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        batch.to(*params["to_args"], **params["to_kwargs"])
        return batch

class MonaiAdapter(Transform):
    """Wrap a MONAI transform for use in TorchIO pipelines.

    Both **dictionary transforms** (subclasses of MONAI's
    `MapTransform`, e.g., `NormalizeIntensityd`) and **array
    transforms** (e.g., `NormalizeIntensity`) are supported.

    Dictionary transforms operate on the full subject dictionary:
    only the keys specified in the MONAI transform are modified.

    Array transforms are applied to each
    [`ScalarImage`][torchio.ScalarImage] in the subject individually,
    respecting the `include` / `exclude` parameters.

    Args:
        monai_transform: A MONAI transform or any callable. Requires
            MONAI to be installed: `pip install torchio[monai]`.
        **kwargs: See [`Transform`][torchio.Transform] for additional
            keyword arguments.

    Examples:
        >>> import torchio as tio
        >>> from monai.transforms import NormalizeIntensity
        >>> # Array transform: applied to each ScalarImage
        >>> adapter = tio.MonaiAdapter(NormalizeIntensity())
        >>> result = adapter(subject)
        >>> # Inside a Compose pipeline
        >>> pipeline = tio.Compose([
        ...     tio.MonaiAdapter(NormalizeIntensity()),
        ...     tio.Noise(std=0.1),
        ... ])

    Note:
        `MonaiAdapter` does **not** record itself in the subject's
        transform history, because MONAI transform objects are not
        serializable.
    """

    def __init__(self, monai_transform: Callable, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        if not callable(monai_transform):
            msg = (
                "monai_transform must be callable, "
                f"got {type(monai_transform).__name__}"
            )
            raise TypeError(msg)
        self.monai_transform = monai_transform

    def forward(self, data):
        """Apply without recording history (MONAI transforms are opaque)."""
        batch, unwrap = self._wrap(data)
        if self.copy:
            batch = _copy.deepcopy(batch)
        if torch.rand(1).item() > self.p:
            return unwrap(batch)
        # MONAI transforms operate per-subject
        monai = get_monai()
        subjects = batch.unbatch()
        for subject in subjects:
            is_dict = isinstance(
                self.monai_transform,
                monai.transforms.MapTransform,
            )
            if is_dict:
                _apply_dict_transform(subject, self.monai_transform, monai)
            else:
                images = self._get_subject_images(subject)
                _apply_array_transform(images, self.monai_transform, monai)
        from ..data.batch import SubjectsBatch

        result = SubjectsBatch.from_subjects(subjects)
        result.applied_transforms = batch.applied_transforms
        return unwrap(result)

    def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any:
        # Not used: MonaiAdapter overrides forward directly
        return batch

    def _get_subject_images(self, subject: Subject) -> dict[str, Image]:
        """Filter to ScalarImage, then apply include/exclude."""
        images: dict[str, Image] = {
            k: v for k, v in subject.images.items() if isinstance(v, ScalarImage)
        }
        if self.include is not None:
            images = {k: v for k, v in images.items() if k in self.include}
        if self.exclude is not None:
            images = {k: v for k, v in images.items() if k not in self.exclude}
        return images