Tutorial proposal: complete AlbumentationsX augmentation pipeline for HoVerNet

[ternaus]

Is your feature request related to a problem? Please describe.

Histopathology models can encounter both spatial variation and H&E color shifts between laboratories and staining protocols. The current HoVerNet training example combines MONAI affine and flip transforms with image-only smoothing, noise, and generic RGB ColorJitter.

MONAI already has an open request for native H&E-space jitter in Project-MONAI/MONAI#5809, with an open implementation PR in Project-MONAI/MONAI#5994. This proposal is for an optional tutorial using a released external augmentation pipeline. The open MONAI work would remain the path for a native transform.

AlbumentationsX 2.4.5 extends HEStain with an explicit full-rank H/E/DAB basis and independent perturbation of the third stain component. It can also apply one sampled affine or flip consistently to an image, instance mask, and type mask while keeping blur, noise, and stain perturbation image-only.

Describe the solution you'd like

Would you be open to a focused tutorial that replaces the current stochastic augmentation block with one AlbumentationsX pipeline while preserving the surrounding MONAI HoVerNet workflow?

The responsibilities would remain explicit:

• MONAI loads the CoNSeP patches, prepares the instance and type targets, performs deterministic output-shape processing, computes HoVer maps, and owns the model, loss, training, and evaluation workflow.
• AlbumentationsX owns the stochastic augmentation policy: affine, horizontal and vertical flips, blur or noise, and HED-space perturbation.
• Affine and flips update image, label_inst, and label_type together. Blur, noise, and HEStain update only image.
• MONAI's random state seeds each AlbumentationsX invocation, so CacheDataset and DataLoader workers continue to treat the adapter as a random transform.

The HoVerNet trainer does not consume physical-space affine metadata after augmentation. The adapter therefore converts the three patch targets to plain tensors at this boundary. A workflow that relies on MetaTensor.affine after random geometry would need to keep that geometry in MONAI or explicitly update the metadata; this tutorial would not claim that unchanged affine metadata remains valid.

A tutorial-specific MapTransform adapter could define the complete stochastic policy in one place:

import albumentations as A
import cv2
import numpy as np

from monai.data import MetaTensor
from monai.transforms import MapTransform, RandomizableTransform
from monai.utils import MAX_SEED


HED_BASIS = np.array(
    [
        [0.65, 0.70, 0.29],  # Hematoxylin
        [0.07, 0.99, 0.11],  # Eosin
        [0.27, 0.57, 0.78],  # DAB
    ],
    dtype=np.float32,
)


class AlbumentationsXHoVerNetd(RandomizableTransform, MapTransform):
    def __init__(self):
        keys = ("image", "label_inst", "label_type")
        MapTransform.__init__(self, keys)
        RandomizableTransform.__init__(self, prob=1.0)
        theta = 0.05
        self.transform = A.Compose(
            [
                A.Affine(
                    scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
                    translate_px={"x": (-6, 6), "y": (-6, 6)},
                    rotate=(-180, 180),
                    shear={"x": (-3, 3), "y": (-3, 3)},
                    interpolation=cv2.INTER_LINEAR,
                    mask_interpolation=cv2.INTER_NEAREST,
                    border_mode=cv2.BORDER_CONSTANT,
                    fill=0,
                    fill_mask=0,
                    keep_ratio=False,
                    p=1.0,
                ),
                A.VerticalFlip(p=0.5),
                A.HorizontalFlip(p=0.5),
                A.OneOf(
                    [
                        A.GaussianBlur(
                            blur_range=(0, 0),
                            sigma_range=(0.1, 1.1),
                            p=1.0,
                        ),
                        A.MedianBlur(blur_range=(3, 3), p=1.0),
                        A.GaussNoise(std_range=(0.01, 0.05), p=1.0),
                    ],
                    p=1.0,
                ),
                A.HEStain(
                    method="custom",
                    stain_matrix=HED_BASIS,
                    residual_mode="augment",
                    intensity_scale_range=(1 - theta, 1 + theta),
                    intensity_shift_range=(-theta, theta),
                    augment_background=True,
                    p=0.5,
                ),
            ],
            additional_targets={
                "label_type": "mask",
            },
            strict=True,
        )

    def __call__(self, data):
        result = dict(data)
        self.randomize(None)
        invocation_seed = int(self.R.randint(MAX_SEED, dtype="uint32"))

        image = result["image"]
        label_inst = result["label_inst"]
        label_type = result["label_type"]
        if isinstance(image, MetaTensor):
            image = image.as_tensor()
        if isinstance(label_inst, MetaTensor):
            label_inst = label_inst.as_tensor()
        if isinstance(label_type, MetaTensor):
            label_type = label_type.as_tensor()

        augmented = self.transform(
            image=image,
            mask=label_inst,
            label_type=label_type,
            invocation_seed=invocation_seed,
        )
        result["image"] = augmented["image"]
        result["label_inst"] = augmented["mask"]
        result["label_type"] = augmented["label_type"]
        return result

In the existing pipeline, MONAI would prepare supported dtypes before the AX boundary. The stochastic MONAI and TorchVision transforms would be replaced by the adapter, while deterministic HoVerNet preprocessing would remain:

LoadImaged(keys=["image", "label_inst", "label_type"], image_only=True),
EnsureChannelFirstd(
    keys=["image", "label_inst", "label_type"],
    channel_dim=-1,
),
Lambdad(keys="label_inst", func=lambda x: measure.label(x)),
CastToTyped(keys="image", dtype=np.uint8),
CastToTyped(keys=["label_inst", "label_type"], dtype=np.int16),
AlbumentationsXHoVerNetd(),

# Existing deterministic HoVerNet preprocessing continues here.
CenterSpatialCropd(keys="image", roi_size=cfg["patch_size"]),
AsDiscreted(keys="label_type", to_onehot=5),
ScaleIntensityRanged(
    keys="image",
    a_min=0,
    a_max=255,
    b_min=0.0,
    b_max=1.0,
    clip=True,
),
# ComputeHoVerMapsd and the existing output-target crops follow.

The tutorial would:

1. visualize the augmented RGB patch with its instance boundaries and type mask;
2. explain which transforms share geometry across all three targets and which transforms affect only the image;
3. verify image-mask shape agreement, nearest-neighbor mask interpolation, instance IDs, and type IDs after augmentation;
4. show how MONAI controls the random state used for each AX invocation; and
5. document the plain-tensor metadata boundary for this patch-training workflow.

This would demonstrate a complete augmentation boundary inside MONAI's pathology data and training workflow. It would make no claim that this policy improves model quality without a benchmark.

The contribution could be either a focused section in the existing HoVerNet material or a small standalone notebook under pathology/, whichever is easier for the maintainers to own.

Describe alternatives you've considered

• Keep the current MONAI and TorchVisiond(ColorJitter) augmentation block. This remains the smallest dependency surface and may be preferable for the existing example.
• Wait for #5994 and keep the complete augmentation policy native to MONAI. That may be preferable if the maintainers want to avoid an optional third-party dependency in the pathology tutorials.
• Add only an HEStain adapter. That would solve the stain-space part, but it would leave stochastic geometry and image-only transforms split across two augmentation systems in the same tutorial.
• Add the example to modules/integrate_3rd_party_transforms.ipynb. That notebook establishes a useful integration precedent, but its current example is a 3D Spleen workflow; the user problem here is specific to 2D pathology and the HoVerNet image/instance/type contract.

Additional context

• The public package is installed as albumentationsx and imported as albumentations.
• AlbumentationsX is AGPL-3.0-only and requires Python 3.10 or newer.
• Importing AlbumentationsX requires PyTorch. The package does not select a PyTorch build because users need to choose CPU, CUDA, or MPS; this example would reuse the same user-selected PyTorch runtime as MONAI.
• AlbumentationsX would be an optional tutorial dependency only, with no change to MONAI's default dependencies.

I tested the adapter against MONAI Tutorials at b3c6c99, MONAI at 605611b, and AlbumentationsX 2.4.5. A synthetic image, instance mask, and type mask retained their shapes and dtypes; the two masks received the same sampled geometry with nearest-neighbor interpolation, and image-only transforms left their values untouched. Re-seeding through MONAI reproduced the same sequence, CacheDataset reapplied the random policy, and two DataLoader workers produced different realizations.

I also searched the current MONAI and MONAI Tutorials code, Issues, pull requests, and Discussions. I found the native work above and the existing third-party integration notebook, but no AlbumentationsX pathology tutorial or duplicate proposal.

If this use case fits the tutorials, I would be happy to implement it after the maintainers confirm the preferred placement and whether they would rather wait for the native transform.