ahcore 0.1.0__py3-none-any.whl

ahcore/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2025 Jonas Teuwen & Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ahcore: reusable building blocks for whole-slide-image deep learning pipelines.
+
+A manifest-backed data catalog, pluggable adapters, a foundation-model feature store, MIL / encoder /
+fusion and segmentation model heads, a Task-driven PyTorch-Lightning module, and Zarr tile readers /
+writers. Import the pieces from their subpackages (e.g. ``ahcore.models``, ``ahcore.data``,
+``ahcore.tasks``, ``ahcore.readers``).
+"""
+
+# Single source of truth for the version; pyproject.toml reads it via hatch's version hook.
+__version__ = "0.1.0"
+
+import ahcore.writers as writers
+
+__all__ = ["writers"]

ahcore/augmentations/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .adapter_augmentation import AdapterAugmentation
+from .interfaces import AdapterAugmentor
+from .tile_features import SelectRandomTiles, TileFeaturesAugmentation
+from .tile_mask_pair import HEDColorAugmentation, TileMaskPairAugmentation
+
+__all__ = [
+    "AdapterAugmentation",
+    "AdapterAugmentor",
+    "HEDColorAugmentation",
+    "SelectRandomTiles",
+    "TileFeaturesAugmentation",
+    "TileMaskPairAugmentation",
+]

ahcore/augmentations/adapter_augmentation.py

@@ -0,0 +1,35 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import torch.nn as nn
+
+from ahcore.augmentations.interfaces import AdapterAugmentor
+from ahcore.data.interfaces import AdapterBundle
+
+
+class AdapterAugmentation(nn.Module):
+    """Composes an ordered list of :class:`AdapterAugmentor` modules into one bundle-to-bundle transform.
+
+    ``forward`` threads the ``AdapterBundle`` through each augmentor in turn and returns the result.
+    """
+
+    def __init__(self, augmentors: list[AdapterAugmentor]) -> None:
+        super().__init__()
+        self._augmentors = nn.ModuleList(augmentors)
+
+    def forward(self, bundle: AdapterBundle) -> AdapterBundle:
+        for augmentor in self._augmentors:
+            bundle = augmentor(bundle)
+        return bundle

ahcore/augmentations/interfaces.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import torch.nn as nn
+
+from ahcore.data.interfaces import AdapterBundle
+
+
+class AdapterAugmentor(nn.Module):
+    """Interface for an augmentor that maps an :class:`AdapterBundle` to an augmented ``AdapterBundle``."""
+
+    def __call__(self, bundle: AdapterBundle) -> AdapterBundle:
+        raise NotImplementedError

ahcore/augmentations/tile_features.py

@@ -0,0 +1,104 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import hashlib
+from typing import Callable
+
+import torch
+import torch.nn as nn
+
+from ahcore.data.interfaces import AdapterBundle
+
+
+class TileFeaturesAugmentation(nn.Module):
+    """Runs an ordered list of bundle-to-bundle transforms over a tile-features :class:`AdapterBundle`.
+
+    Each transform maps an ``AdapterBundle`` to an ``AdapterBundle``; ``forward`` applies them in order.
+    Use :meth:`for_classification` for the default tile-subsampling pipeline.
+    """
+
+    def __init__(self, transforms: list[Callable[[AdapterBundle], AdapterBundle]]) -> None:
+        super().__init__()
+        self._transforms = transforms
+
+    @classmethod
+    def for_classification(
+        cls, num_tiles: int = 1000, seed: int | None = None, per_sample: bool = False, meta_key: str = "slide_id"
+    ) -> TileFeaturesAugmentation:
+        transforms: list[Callable[[AdapterBundle], AdapterBundle]] = [
+            SelectRandomTiles(num_tiles=num_tiles, seed=seed, per_sample=per_sample, meta_key=meta_key),
+        ]
+        return TileFeaturesAugmentation(transforms)
+
+    def forward(self, data: AdapterBundle) -> AdapterBundle:
+        for transform in self._transforms:
+            data = transform(data)
+        return data
+
+
+class SelectRandomTiles(nn.Module):
+    """Randomly subsample a fixed number of tiles from a bundle, keeping all per-tile tensors aligned.
+
+    Samples ``num_tiles`` rows without replacement (with replacement if the bag is smaller) and applies
+    the same indices to every per-tile tensor in ``bundle.data`` so features/coords stay row-aligned.
+    With ``seed`` set the choice is deterministic; ``per_sample=True`` derives the seed from
+    ``bundle.meta[meta_key]`` so each slide gets its own reproducible selection.
+    """
+
+    def __init__(
+        self, num_tiles: int, seed: int | None = None, per_sample: bool = False, meta_key: str = "slide_id"
+    ) -> None:
+        super().__init__()
+        self._num_tiles = num_tiles
+        self._seed = seed
+        self._per_sample = per_sample
+        self._meta_key = meta_key
+
+    def _make_generator(self, bundle: AdapterBundle, device: torch.device | str) -> torch.Generator | None:
+        if self._seed is None:
+            return None
+        derived_seed = self._seed
+        if (
+            self._per_sample
+            and hasattr(bundle, "meta")
+            and isinstance(bundle.meta, dict)
+            and self._meta_key in bundle.meta
+        ):
+            key = str(bundle.meta[self._meta_key])
+            digest = hashlib.sha256(f"{self._seed}:{key}".encode("utf-8")).digest()
+            derived_seed = int.from_bytes(digest[:8], "little", signed=False)
+        gen = torch.Generator(device=device)
+        gen.manual_seed(derived_seed)
+        return gen
+
+    def forward(self, bundle: AdapterBundle) -> AdapterBundle:
+        features = bundle.data["features"]
+        num_features = features.shape[0]
+
+        gen = self._make_generator(bundle, features.device)
+
+        if num_features < self._num_tiles:
+            # If there are fewer features than num_tiles, we sample with replacement.
+            indices = torch.randint(0, num_features, (self._num_tiles,), device=features.device, generator=gen)
+        else:
+            # Otherwise, we sample without replacement.
+            indices = torch.randperm(num_features, device=features.device, generator=gen)[: self._num_tiles]
+
+        # Subsample every per-tile tensor by the same indices so features / coords / (future)
+        # patch_tokens stay row-aligned — no per-key hardcoding.
+        for key, value in bundle.data.items():
+            if isinstance(value, torch.Tensor) and value.shape[0] == num_features:
+                bundle.data[key] = value[indices]
+        return bundle

ahcore/augmentations/tile_mask_pair.py

@@ -0,0 +1,255 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import Any, Optional, cast
+
+import kornia.augmentation as K
+import torch
+import torch.nn as nn
+from kornia.augmentation import random_generator as rg
+from kornia.constants import DataKey, Resample
+from omegaconf import ListConfig
+
+from ahcore.data.interfaces import AdapterBundle
+
+
+class TileMaskPairAugmentation(nn.Module):
+    """Kornia augmentation for a segmentation (image, mask, ignore_mask) triple in an ``AdapterBundle``.
+
+    Applies intensity augmentations to the image only, then geometric augmentations jointly to image
+    and masks (nearest-neighbour resampling on the masks), and finally centre-crops all three to
+    ``tile_size``. Use :meth:`for_segmentation` for the default augmentation set.
+    """
+
+    def __init__(
+        self,
+        tile_size: tuple[int, int],
+        intensity_augmentations: list[K.IntensityAugmentationBase2D],
+        geometric_augmentations: list[K.GeometricAugmentationBase2D],
+    ) -> None:
+        super().__init__()
+        tile_size = (int(tile_size[0]), int(tile_size[1]))  # Handle omegaconf ListConfig
+        self._intensity_augmentations = K.AugmentationSequential(
+            *intensity_augmentations,
+            data_keys=[DataKey.INPUT],
+            same_on_batch=False,
+            random_apply=1,
+        )
+
+        self._geometric_augmentations = K.AugmentationSequential(
+            *geometric_augmentations,
+            data_keys=[DataKey.INPUT, DataKey.MASK, DataKey.MASK],
+            same_on_batch=False,
+            extra_args={DataKey.MASK: dict(resample=Resample.NEAREST, align_corners=True)},
+            random_apply=1,
+        )
+
+        self._final_transforms = K.AugmentationSequential(
+            K.CenterCrop(size=tile_size),
+            data_keys=[DataKey.INPUT, DataKey.MASK, DataKey.MASK],
+            same_on_batch=False,
+            extra_args={DataKey.MASK: dict(resample=Resample.NEAREST, align_corners=True)},
+        )
+
+    @classmethod
+    def for_segmentation(
+        cls,
+        tile_size: tuple[int, int] = (1024, 1024),
+    ) -> TileMaskPairAugmentation:
+        intensity_augmentations = [
+            # HEDColorAugmentation(p=0.5, scale_sigma=0.2, bias_sigma=0.2),
+            K.ColorJitter(p=0.5, brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05),
+            K.RandomGaussianBlur(p=0.5, kernel_size=(9, 9), sigma=(0.1, 1.0)),
+            K.RandomSharpness(p=0.5, sharpness=10),
+        ]
+        geometric_augmentations = [
+            K.RandomHorizontalFlip(p=0.5),
+            K.RandomVerticalFlip(p=0.5),
+            K.RandomPerspective(p=0.5, distortion_scale=0.5),
+            K.RandomAffine(p=0.5, degrees=(0, 90)),
+            K.RandomAffine(p=0.5, degrees=0, scale=(0.75, 1.0)),
+            K.RandomAffine(p=0.5, degrees=0, translate=(0, 0.05)),
+            K.RandomAffine(p=0.5, degrees=0, shear=(0, 5)),
+        ]
+        return TileMaskPairAugmentation(tile_size, intensity_augmentations, geometric_augmentations)
+
+    @staticmethod
+    def _to_kornia(mask: torch.Tensor) -> torch.Tensor:
+        """Present a mask to kornia as a float [B, C, H, W]: grid_sample needs floats, and a bool input
+        would otherwise demote the co-warped masks. A 2D [H, W] label map gains explicit batch/channel dims."""
+        while mask.ndim < 4:
+            mask = mask.unsqueeze(0)
+        return mask.float()
+
+    @staticmethod
+    def _from_kornia(mask: torch.Tensor, ndim: int, dtype: torch.dtype) -> torch.Tensor:
+        """Undo :meth:`_to_kornia`: drop the added batch (and, for a 2D input, channel) dims and restore
+        the original dtype. NEAREST resampling keeps values exact, so casting back is lossless."""
+        mask = mask.squeeze(0)  # drop batch
+        if ndim == 2:
+            mask = mask.squeeze(0)  # drop the channel added for a 2D label map
+        return mask.to(dtype)
+
+    def forward(self, bundle: AdapterBundle) -> AdapterBundle:
+        image = bundle.data["image"].unsqueeze(0)  # kornia augmentations expect a batch dimension
+        mask, ignore_mask = bundle.data["mask"], bundle.data["ignore_mask"]
+
+        mask_ndim, mask_dtype = mask.ndim, mask.dtype
+        ignore_ndim, ignore_dtype = ignore_mask.ndim, ignore_mask.dtype
+        mask = self._to_kornia(mask)
+        ignore_mask = self._to_kornia(ignore_mask)
+
+        image = self._intensity_augmentations(image)
+        image, mask, ignore_mask = self._geometric_augmentations(image, mask, ignore_mask)
+        image, mask, ignore_mask = self._final_transforms(image, mask, ignore_mask)
+
+        bundle.data["image"] = image.squeeze(0)
+        bundle.data["mask"] = self._from_kornia(mask, mask_ndim, mask_dtype)
+        bundle.data["ignore_mask"] = self._from_kornia(ignore_mask, ignore_ndim, ignore_dtype)
+        return bundle
+
+
+class HEDColorAugmentation(K.IntensityAugmentationBase2D):
+    """
+    A torch implementation of the color stain augmentation algorithm on the
+    deconvolved Hemaetoxylin-Eosin-DAB (HED) channels of an image as described
+    by Tellez et al. (2018) in Appendix A & B here: https://arxiv.org/pdf/1808.05896.pdf.
+    """
+
+    # Normalized OD matrix from Ruifrok et al. (2001)
+    HED_REFERENCE = torch.Tensor([[0.65, 0.70, 0.29], [0.07, 0.99, 0.11], [0.27, 0.57, 0.78]])
+
+    def __init__(
+        self,
+        scale_sigma: float | list[float] | ListConfig,
+        bias_sigma: float | list[float] | ListConfig,
+        epsilon: float = 1e-6,
+        clamp_output_range: Optional[tuple[float, float]] = (0, 1),
+        p: float = 0.5,
+        p_batch: float = 1.0,
+        same_on_batch: bool = False,
+        keepdim: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Apply a color stain augmentation in the Hemaetoxylin-Eosin-DAB (HED) color space based on [1].
+        The fixed normalized OD matrix values are based on [2].
+
+        Parameters
+        ----------
+        scale_sigma: float, ListConfig or list of floats
+            For each channel in the HED space a random scaling factor is drawn from alpha_i ~ U(1-sigma_i,1+sigma_i).
+        bias_sigma: float, ListConfig or list of floats
+            For each channel in the HED space a random bias is added drawn from beta_i ~ U(-sigma_i,sigma_i).
+        epsilon: float
+            Small positive bias to avoid numerical errors
+        clamp_output_range: tuple of floats or None
+            Clamp output in range after augmenting the input. Conventionally a range of [0, 1] is used, but this will
+            discard some color information. `scale_sigma` and `bias_sigma` should be chosen carefully to prevent this.
+
+        NOTE: To reproduce augmentations as shown in [3], a larger positive bias of `epsilon=2` in combination with
+        `clamp_output_range=(0, 1)` should be used. Then, use `scale_sigma=bias_sigma=0.05` for `HED-light` and
+        `scale_sigma=bias_sigma=0.2` for `HED-strong`.
+        When using the default `epsilon=1e-6` and `clamp_output_range=(0, 1)`, this is not equal but comparable to
+        `scale_sigma=bias_sigma=0.15` for `HED-light` and `scale_sigma=bias_sigma=0.8` for `HED-strong`?
+
+
+        References
+        ----------
+        [1] Tellez, David, et al. "Whole-slide mitosis detection in H&E breast histology using PHH3 as
+            a reference to train distilled stain-invariant convolutional networks."
+            IEEE transactions on medical imaging 37.9 (2018): 2126-2136.
+        [2] Ruifrok AC, Johnston DA. Quantification of histochemical staining by color deconvolution.
+            Anal Quant Cytol Histol. 2001 Aug;23(4):291-9. PMID: 11531144.
+        [3] Tellez, David, et al. "Quantifying the effects of data augmentation and stain color normalization
+            in convolutional neural networks for computational pathology."
+            Medical image analysis 58 (2019): 101544.
+        """
+        super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim)
+
+        if isinstance(scale_sigma, ListConfig):
+            scale_sigma = list(scale_sigma)
+        if isinstance(bias_sigma, ListConfig):
+            bias_sigma = list(bias_sigma)
+
+        # Accept int too: an int is a valid `float` argument, and a whole-number config value (e.g. YAML
+        # `0` or `1`) arrives as int, so a scalar must broadcast whether it comes in as int or float.
+        if isinstance(scale_sigma, (int, float)):
+            scale_sigma = [scale_sigma] * 3
+
+        if isinstance(bias_sigma, (int, float)):
+            bias_sigma = [bias_sigma] * 3
+
+        if len(scale_sigma) not in [1, 3] or len(bias_sigma) not in [1, 3]:
+            raise ValueError(
+                f"scale_sigma and bias_sigma should have either 1 or 3 values, "
+                f"got {scale_sigma} and {bias_sigma} instead."
+            )
+
+        assert isinstance(scale_sigma, list)
+        assert isinstance(bias_sigma, list)
+
+        _scale_sigma = torch.Tensor(scale_sigma).float()
+        _bias_sigma = torch.Tensor(bias_sigma).float()
+
+        scale_factor = torch.stack([1.0 - _scale_sigma, 1.0 + _scale_sigma], dim=0)
+        bias_factor = torch.stack([-_bias_sigma, _bias_sigma], dim=0)  # pylint:disable=E1130
+
+        self._param_generator = rg.PlainUniformGenerator(
+            (scale_factor, "scale", None, None), (bias_factor, "bias", None, None)
+        )
+        self.flags = {
+            "epsilon": torch.tensor([epsilon]),
+            "M": self.HED_REFERENCE,
+            "M_inv": torch.linalg.inv(self.HED_REFERENCE),  # pylint:disable=E1102
+            "clamp_output_range": clamp_output_range,
+        }
+
+    def apply_transform(
+        self,
+        input: torch.Tensor,
+        params: dict[str, torch.Tensor],
+        flags: dict[str, Any],
+        transform: Optional[torch.Tensor] = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """
+        Apply HED color augmentation on an input tensor.
+        """
+        assert flags, "Flags should be provided"
+        assert params, "Params should be provided"
+
+        epsilon = flags["epsilon"].to(input)
+        reference_matrix = flags["M"].to(input)
+        reference_matrix_inv = flags["M_inv"].to(input)
+        alpha = params["scale"][:, None, None, :].to(input)
+        beta = params["bias"][:, None, None, :].to(input)
+
+        rgb_tensor = input.permute(0, 2, 3, 1)
+        optical_density = -torch.log(rgb_tensor + epsilon)
+        # Mypy doesn't understand this is a tensor.
+        hed_tensor = cast(torch.Tensor, optical_density @ reference_matrix_inv)
+
+        augmented_hed_tensor = alpha * hed_tensor + beta
+        # Same problem that mypy doesn't understand
+        augmented_rgb_tensor = torch.exp(-augmented_hed_tensor @ reference_matrix) - epsilon
+        augmented_sample = augmented_rgb_tensor.permute(0, 3, 1, 2)
+
+        # The linter seems to require this
+        assert isinstance(augmented_sample, torch.Tensor)
+
+        if flags["clamp_output_range"] is not None:
+            augmented_sample = torch.clamp(augmented_sample, *flags["clamp_output_range"])
+        return augmented_sample

ahcore/data/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .interfaces import (
+    AdapterBundle,
+    ShapePolicy,
+    SlideView,
+    TileView,
+)
+
+__all__ = [
+    "AdapterBundle",
+    "ShapePolicy",
+    "SlideView",
+    "TileView",
+]

ahcore/data/adapters/__init__.py

@@ -0,0 +1,33 @@
+# Copyright 2025 Joren Brunekreef. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Protocol
+
+from ahcore.data import SlideView, TileView
+from ahcore.data.interfaces import AdapterBundle
+
+
+class SlideAdapter(Protocol):
+    """Protocol for an adapter that produces a representation bundle for a slide."""
+
+    def __init__(self) -> None: ...
+
+    def __call__(self, slide_view: SlideView) -> AdapterBundle: ...
+
+
+class TileAdapter(Protocol):
+    """Protocol for an adapter that produces a representation bundle for a tile."""
+
+    def __init__(self) -> None: ...
+
+    def __call__(self, tile_view: TileView) -> AdapterBundle: ...