jolia_atlas_transform.py

# Vendored verbatim from the internal `raidium.rd.models` library for the
# self-contained Hugging Face release. Only imports were rewritten (raidium
# hub base classes -> jolia_shim; sibling modules -> jolia_* names).
# Do not edit by hand: regenerate with scripts/build_hf_jolia.py.

from __future__ import annotations

import copy
from typing import TYPE_CHECKING, Sequence, Union

import torch
import torch.nn.functional as F
try:
    from .jolia_windowing import batch_apply_windowing_vectorized
except ImportError:
    from jolia_windowing import batch_apply_windowing_vectorized

if TYPE_CHECKING:
    import numpy as np

Tensorable = Union[torch.Tensor, "np.ndarray"]


class PrepareVolume:
    """Convert raw volume to float tensor in (D, H, W) order.

    Handles squeeze of trivial leading dim, permute when depth_last, and depth flip.
    """

    def __init__(self, depth_last: bool = True, flip_depth: bool = True) -> None:
        self.depth_last = depth_last
        self.flip_depth = flip_depth

    def __call__(self, volume: Tensorable, **_kwargs: object) -> torch.Tensor:
        vol = torch.as_tensor(volume).float()
        if vol.ndim == 4 and vol.shape[0] == 1:
            vol = vol.squeeze(0)
        if vol.ndim != 3:
            raise ValueError(f"Expected a 3D volume, got shape {vol.shape}")
        if self.depth_last:
            vol = vol.permute(2, 0, 1).contiguous()
        if self.flip_depth:
            vol = torch.flip(vol, dims=[0])
        return vol


class Resample3D:
    """Resample a 3D tensor to match target spacing.

    Use ``mode="trilinear"`` for images (default) and ``mode="nearest"`` for masks.
    """

    def __init__(
        self,
        target_spacing: tuple[float, ...] = (1.0, 1.0, 1.0),
        mode: str = "trilinear",
    ) -> None:
        self.target_spacing = target_spacing
        self.mode = mode

    def __call__(
        self, tensor: torch.Tensor, current_spacing: tuple[float, ...] | None = None, **_kwargs: object
    ) -> torch.Tensor:
        if current_spacing is None:
            raise ValueError("Resample3D requires current_spacing (from metadata)")
        x = tensor.unsqueeze(0).unsqueeze(0)
        original_shape = x.shape[2:]
        scaling_factors = [current_spacing[i] / self.target_spacing[i] for i in range(len(original_shape))]
        new_shape = [max(int(original_shape[i] * scaling_factors[i]), 1) for i in range(len(original_shape))]
        interp_kwargs: dict[str, object] = {"size": new_shape, "mode": self.mode}
        if self.mode in {"linear", "bilinear", "trilinear", "bicubic"}:
            interp_kwargs["align_corners"] = False
        resized = F.interpolate(x, **interp_kwargs)
        return resized.squeeze(0).squeeze(0)


class Dilate3D:
    """3D morphological dilation via max pooling.

    Inserted into the mask pipeline before ``Resample3D`` to thicken sparse labels so
    they survive aggressive nearest-neighbor downsampling that would otherwise drop them.
    Implemented as ``F.max_pool3d`` with stride 1 and ``kernel_size // 2`` zero padding.

    For multi-class masks, max-pooling propagates the higher label into neighboring voxels
    at label boundaries — acceptable for binary masks; for multi-class with adjacent labels
    expect a one-voxel boundary bias toward the higher label.
    Handles 3D (D, H, W), 4D (B, D, H, W), and 5D (B, C, D, H, W) tensors.
    """

    def __init__(self, kernel_size: int = 3) -> None:
        if kernel_size < 1 or kernel_size % 2 == 0:
            raise ValueError(f"kernel_size must be a positive odd integer, got {kernel_size}")
        self.kernel_size = kernel_size

    def __call__(self, tensor: torch.Tensor, **_kwargs: object) -> torch.Tensor:
        orig_ndim = tensor.dim()
        x = tensor
        if orig_ndim == 3:
            x = x.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 4:
            x = x.unsqueeze(1)
        original_dtype = x.dtype
        pad = self.kernel_size // 2
        dilated = F.max_pool3d(x.float(), kernel_size=self.kernel_size, stride=1, padding=pad)
        dilated = dilated.to(dtype=original_dtype)
        if orig_ndim == 3:
            dilated = dilated.squeeze(0).squeeze(0)
        elif orig_ndim == 4:
            dilated = dilated.squeeze(1)
        return dilated


class Crop3D:
    """Device-agnostic 3D crop.

    H and W are center-cropped, depth is random in training and centered in eval.
    Per-axis ``d_start``/``h_start``/``w_start`` may be passed via kwargs to override
    the default center/random offset (used by ``AtlasImageMaskTransform`` for
    mask-centered cropping).
    Handles 3D (D, H, W), 4D (B, D, H, W) and 5D (B, C, D, H, W) tensors.
    """

    def __init__(
        self,
        target_shape: tuple[int, int, int] = (192, 192, 192),
        training: bool = True,
    ) -> None:
        self.target_shape = target_shape
        self.training = training

    def __call__(
        self,
        images: torch.Tensor,
        d_start: int | None = None,
        h_start: int | None = None,
        w_start: int | None = None,
        **_kwargs: object,
    ) -> torch.Tensor:
        orig_ndim = images.dim()
        if orig_ndim == 3:
            images = images.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 4:
            images = images.unsqueeze(1)

        _, _, d, h, w = images.shape
        td, th, tw = self.target_shape

        if h_start is None:
            h_start = max((h - th) // 2, 0)
        if w_start is None:
            w_start = max((w - tw) // 2, 0)
        if d_start is None:
            if td >= d:
                d_start = 0
            elif self.training:
                d_start = int(torch.randint(0, max(d - td, 1) + 1, ()).item())
            else:
                d_start = (d - td) // 2

        cropped = images[
            :,
            :,
            d_start : d_start + min(td, d),
            h_start : h_start + min(th, h),
            w_start : w_start + min(tw, w),
        ]

        if orig_ndim == 3:
            cropped = cropped.squeeze(0).squeeze(0)
        elif orig_ndim == 4:
            cropped = cropped.squeeze(1)
        return cropped


class Pad3D:
    """Device-agnostic 3D center-pad to target shape.

    Handles 3D (D, H, W), 4D (B, D, H, W) and 5D (B, C, D, H, W) tensors.
    """

    def __init__(
        self,
        target_shape: tuple[int, int, int] = (192, 192, 192),
        padding_value: float = -1024.0,
    ) -> None:
        self.target_shape = target_shape
        self.padding_value = padding_value

    def __call__(self, images: torch.Tensor, **_kwargs: object) -> torch.Tensor:
        orig_ndim = images.dim()
        if orig_ndim == 3:
            images = images.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 4:
            images = images.unsqueeze(1)

        _, _, cd, ch, cw = images.shape
        td, th, tw = self.target_shape

        if cd < td or ch < th or cw < tw:
            diff_d = max(td - cd, 0)
            diff_h = max(th - ch, 0)
            diff_w = max(tw - cw, 0)
            images = F.pad(
                images,
                (
                    diff_w // 2,
                    diff_w - diff_w // 2,
                    diff_h // 2,
                    diff_h - diff_h // 2,
                    diff_d // 2,
                    diff_d - diff_d // 2,
                ),
                value=self.padding_value,
            )

        if orig_ndim == 3:
            images = images.squeeze(0).squeeze(0)
        elif orig_ndim == 4:
            images = images.squeeze(1)
        return images


class Rotation3D:
    """Device-agnostic random rotation using affine_grid + grid_sample.

    Supports rotation around Z-axis only (default) or all three axes.
    Operates in float32 for numerical stability, casts back to original dtype.
    Handles 3D (D, H, W), 4D (B, D, H, W) and 5D (B, C, D, H, W) tensors.
    """

    def __init__(
        self,
        degrees: float = 10.0,
        p: float = 0.5,
        training: bool = True,
        axes: str = "z",
    ) -> None:
        self.degrees = degrees
        self.p = p
        self.training = training
        self.axes = axes

    @staticmethod
    def _random_angle(degrees: float) -> float:
        return (torch.rand(1).item() * 2 - 1) * degrees

    @staticmethod
    def _rot_z(angle_rad: float, device: torch.device) -> torch.Tensor:
        c = float(torch.cos(torch.tensor(angle_rad)))
        s = float(torch.sin(torch.tensor(angle_rad)))
        return torch.tensor([[c, -s, 0], [s, c, 0], [0, 0, 1]], dtype=torch.float32, device=device)

    @staticmethod
    def _rot_x(angle_rad: float, device: torch.device) -> torch.Tensor:
        c = float(torch.cos(torch.tensor(angle_rad)))
        s = float(torch.sin(torch.tensor(angle_rad)))
        return torch.tensor([[1, 0, 0], [0, c, -s], [0, s, c]], dtype=torch.float32, device=device)

    @staticmethod
    def _rot_y(angle_rad: float, device: torch.device) -> torch.Tensor:
        c = float(torch.cos(torch.tensor(angle_rad)))
        s = float(torch.sin(torch.tensor(angle_rad)))
        return torch.tensor([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=torch.float32, device=device)

    def __call__(self, images: torch.Tensor, **_kwargs: object) -> torch.Tensor:
        if not self.training or torch.rand(1).item() > self.p:
            return images

        orig_ndim = images.dim()
        if orig_ndim == 3:
            images = images.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 4:
            images = images.unsqueeze(1)

        original_dtype = images.dtype
        images_f32 = images.float()
        b, c, d, h, w = images_f32.shape
        device = images_f32.device
        pi = 3.141592653589793

        angle_z = self._random_angle(self.degrees) * pi / 180.0
        rot_mat = self._rot_z(angle_z, device)

        if self.axes == "all":
            angle_x = self._random_angle(self.degrees) * pi / 180.0
            angle_y = self._random_angle(self.degrees) * pi / 180.0
            rot_mat = rot_mat @ self._rot_x(angle_x, device) @ self._rot_y(angle_y, device)

        affine = torch.zeros(b, 3, 4, dtype=torch.float32, device=device)
        affine[:, :3, :3] = rot_mat.unsqueeze(0)

        grid = F.affine_grid(affine, [b, c, d, h, w], align_corners=False)
        rotated = F.grid_sample(images_f32, grid, mode="bilinear", padding_mode="border", align_corners=False)
        rotated = rotated.to(dtype=original_dtype)

        if orig_ndim == 3:
            rotated = rotated.squeeze(0).squeeze(0)
        elif orig_ndim == 4:
            rotated = rotated.squeeze(1)
        return rotated


class RandomIntensityShift:
    """Random intensity scale and shift for Hounsfield unit augmentation.

    Applied before windowing on raw HU values.
    Handles 3D (D, H, W), 4D (B, D, H, W) and 5D (B, C, D, H, W) tensors.
    """

    def __init__(
        self,
        scale_range: tuple[float, float] = (0.95, 1.05),
        shift_range: tuple[float, float] = (-10.0, 10.0),
        p: float = 0.5,
        training: bool = True,
    ) -> None:
        self.scale_range = scale_range
        self.shift_range = shift_range
        self.p = p
        self.training = training

    def __call__(self, images: torch.Tensor, **_kwargs: object) -> torch.Tensor:
        if not self.training or torch.rand(1).item() > self.p:
            return images
        scale = torch.empty(1, device=images.device).uniform_(self.scale_range[0], self.scale_range[1]).item()
        shift = torch.empty(1, device=images.device).uniform_(self.shift_range[0], self.shift_range[1]).item()
        return images * scale + shift


class RandomCrop3D:
    """Random crop with jitter on all three axes.

    Unlike Crop3D which center-crops H/W, this applies a random offset on all axes
    during training for spatial diversity. Per-axis ``d_start``/``h_start``/``w_start``
    may be passed via kwargs to override the default random/center offset (used by
    ``AtlasImageMaskTransform`` for mask-centered cropping).
    Handles 3D (D, H, W), 4D (B, D, H, W) and 5D (B, C, D, H, W) tensors.
    """

    def __init__(
        self,
        target_shape: tuple[int, int, int] = (192, 192, 192),
        training: bool = True,
    ) -> None:
        self.target_shape = target_shape
        self.training = training

    def __call__(
        self,
        images: torch.Tensor,
        d_start: int | None = None,
        h_start: int | None = None,
        w_start: int | None = None,
        **_kwargs: object,
    ) -> torch.Tensor:
        orig_ndim = images.dim()
        if orig_ndim == 3:
            images = images.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 4:
            images = images.unsqueeze(1)

        _, _, d, h, w = images.shape
        td, th, tw = self.target_shape

        if d_start is None:
            if self.training:
                d_start = int(torch.randint(0, max(d - td, 1) + 1, ()).item()) if d > td else 0
            else:
                d_start = max((d - td) // 2, 0)
        if h_start is None:
            if self.training:
                h_start = int(torch.randint(0, max(h - th, 1) + 1, ()).item()) if h > th else 0
            else:
                h_start = max((h - th) // 2, 0)
        if w_start is None:
            if self.training:
                w_start = int(torch.randint(0, max(w - tw, 1) + 1, ()).item()) if w > tw else 0
            else:
                w_start = max((w - tw) // 2, 0)

        cropped = images[
            :,
            :,
            d_start : d_start + min(td, d),
            h_start : h_start + min(th, h),
            w_start : w_start + min(tw, w),
        ]

        if orig_ndim == 3:
            cropped = cropped.squeeze(0).squeeze(0)
        elif orig_ndim == 4:
            cropped = cropped.squeeze(1)
        return cropped


class ApplyWindowing:
    """Vectorized CT windowing.

    Handles 3D (D, H, W), 4D (B, D, H, W) and 5D (B, C, D, H, W) input.
    The channel dimension is expanded by windowing (1 -> N windows).
    """

    def __init__(
        self,
        window_type: str | list[str] = "all",
        modality: str = "CT",
        dtype: torch.dtype = torch.bfloat16,
    ) -> None:
        self.window_type = window_type
        self.modality = modality
        self.dtype = dtype

    def __call__(self, images: torch.Tensor, **_kwargs: object) -> torch.Tensor:
        orig_ndim = images.dim()
        if orig_ndim == 3:
            images = images.unsqueeze(0).unsqueeze(0)
        elif orig_ndim == 4:
            images = images.unsqueeze(1)

        result = batch_apply_windowing_vectorized(
            images,
            windows=self.window_type,
            modality=self.modality,
            torch_operating_dtype=self.dtype,
        )

        if orig_ndim == 3:
            result = result.squeeze(0)
        return result


def _reorder_hwd_to_dhw(t: tuple) -> tuple:
    """Reorder a (H, W, D) tuple to (D, H, W)."""
    return (t[2], t[0], t[1])


def _offsets_to_kwargs(offsets: tuple[int, int, int] | None) -> dict[str, int]:
    if offsets is None:
        return {}
    return {"d_start": offsets[0], "h_start": offsets[1], "w_start": offsets[2]}


def find_mask_centered_offsets(mask: torch.Tensor, target_shape: tuple[int, int, int]) -> tuple[int, int, int] | None:
    """Compute ``(d_start, h_start, w_start)`` so a ``target_shape`` crop is centered on
    the mask centroid, clamped to volume bounds.

    Operates on the last three spatial dims of ``mask``. Returns ``None`` if the mask is
    empty (all zero).
    """
    spatial = mask if mask.dim() == 3 else mask.reshape(-1, *mask.shape[-3:]).any(dim=0).float()
    indices = (spatial != 0).nonzero(as_tuple=False).float()
    if indices.numel() == 0:
        return None
    centroid = indices.mean(dim=0)
    d_dim, h_dim, w_dim = spatial.shape[-3], spatial.shape[-2], spatial.shape[-1]
    td, th, tw = target_shape

    def _clamp(c: float, dim: int, target: int) -> int:
        max_start = max(0, dim - target)
        return max(0, min(int(c - target / 2), max_start))

    return (
        _clamp(centroid[0].item(), d_dim, td),
        _clamp(centroid[1].item(), h_dim, th),
        _clamp(centroid[2].item(), w_dim, tw),
    )


class AtlasTransform:
    transform_input = "metadata"

    def __init__(
        self,
        precomputed: bool = False,
        depth_last: bool = True,
        training: bool = True,
        cpu_transforms: Sequence | None = None,
        **_kwargs: object,
    ) -> None:
        self.precomputed = precomputed
        self.depth_last = depth_last

        if cpu_transforms:
            for t in cpu_transforms:
                if hasattr(t, "training"):
                    t.training = training
            self._cpu_transforms = list(cpu_transforms)
        else:
            self._cpu_transforms = []

    def __call__(self, volume: Tensorable, metadata: dict, **kwargs: object) -> torch.Tensor:
        if self.precomputed:
            vol: torch.Tensor | Tensorable = torch.as_tensor(volume).float()
            return vol

        vol = volume
        context = self._build_context(metadata)

        for t in self._cpu_transforms:
            vol = t(vol, **context)

        return vol

    def _build_context(self, metadata: dict | None) -> dict[str, object]:
        context: dict[str, object] = {}
        resolution = metadata.get("resolution") if metadata else None
        if resolution is not None:
            current_spacing = tuple(resolution)
            if self.depth_last:
                current_spacing = _reorder_hwd_to_dhw(current_spacing)
            context["current_spacing"] = current_spacing
        return context

    def run_capturing_crop_offsets(
        self, volume: Tensorable, metadata: dict | None
    ) -> tuple[torch.Tensor, list[tuple[int, int, int] | None]]:
        """Run the pipeline; at each ``Crop3D``/``RandomCrop3D``, compute mask-centered
        offsets from the current volume state and use them. Return the output and the
        list of captured offsets (one per crop, in pipeline order).
        """
        if self.precomputed:
            return torch.as_tensor(volume).float(), []
        vol: object = volume
        context = self._build_context(metadata)
        captured: list[tuple[int, int, int] | None] = []
        for t in self._cpu_transforms:
            if isinstance(t, (Crop3D, RandomCrop3D)):
                offsets = find_mask_centered_offsets(torch.as_tensor(vol), t.target_shape)
                captured.append(offsets)
                vol = t(vol, **context, **_offsets_to_kwargs(offsets))
            else:
                vol = t(vol, **context)
        return vol, captured  # type: ignore[return-value]

    def run_with_crop_offsets(
        self,
        volume: Tensorable,
        metadata: dict | None,
        crop_offsets: list[tuple[int, int, int] | None],
    ) -> torch.Tensor:
        """Run the pipeline; at each ``Crop3D``/``RandomCrop3D``, use the next entry from
        ``crop_offsets`` (consumed in pipeline order). When an entry is ``None`` the
        transform falls back to its default offset logic.
        """
        if self.precomputed:
            return torch.as_tensor(volume).float()
        vol: object = volume
        context = self._build_context(metadata)
        crop_idx = 0
        for t in self._cpu_transforms:
            if isinstance(t, (Crop3D, RandomCrop3D)):
                offsets = crop_offsets[crop_idx] if crop_idx < len(crop_offsets) else None
                crop_idx += 1
                vol = t(vol, **context, **_offsets_to_kwargs(offsets))
            else:
                vol = t(vol, **context)
        return vol  # type: ignore[return-value]

    @classmethod
    def for_mask(cls, image_transform: "AtlasTransform", dilate_kernel: int = 3) -> "AtlasTransform":
        """Build a mask companion mirroring the spatial pipeline of ``image_transform``.

        The mask pipeline:
          - inserts a ``Dilate3D(kernel_size=dilate_kernel)`` before each ``Resample3D``
            so sparse labels survive aggressive nearest-neighbor downsampling
            (set ``dilate_kernel=0`` to disable),
          - replaces ``Resample3D`` with a nearest-neighbor variant,
          - forces ``Pad3D.padding_value=0``,
          - drops intensity-only transforms (``RandomIntensityShift``, ``ApplyWindowing``),
          - is always non-precomputed (masks are not precomputed offline),
          - is always ``training=False`` (deterministic spatial transforms required for image/mask alignment).
        """
        mask_cpu_transforms: list = []
        for t in image_transform._cpu_transforms:
            if isinstance(t, (RandomIntensityShift, ApplyWindowing)):
                continue
            if isinstance(t, Resample3D):
                if dilate_kernel:
                    mask_cpu_transforms.append(Dilate3D(kernel_size=dilate_kernel))
                mask_cpu_transforms.append(Resample3D(target_spacing=t.target_spacing, mode="nearest"))
                continue
            if isinstance(t, Pad3D):
                mask_cpu_transforms.append(Pad3D(target_shape=t.target_shape, padding_value=0))
                continue
            mask_cpu_transforms.append(copy.deepcopy(t))
        return cls(
            precomputed=False,
            depth_last=image_transform.depth_last,
            training=False,
            cpu_transforms=mask_cpu_transforms,
        )