NEW: Add stainaugment feature for stain augmentation

requirements.conda.yml

@@ -6,6 +6,7 @@ channels:
   - defaults
 dependencies:
   - python=3.7
+  - albumentations
   - Click
   - cython
   - defusedxml

requirements.dev.conda.yml

@@ -5,6 +5,7 @@ channels:
   - pytorch
   - defaults
 dependencies:
+  - albumentations
   - black
   - Click
   - cython

requirements.txt

@@ -1,3 +1,4 @@
+albumentations>0.5.0
 Click>=7.0
 defusedxml
 glymur

requirements.win64.conda.yml

@@ -6,6 +6,7 @@ channels:
   - defaults
 dependencies:
   - python=3.7
+  - albumentations
   - Click
   - cython
   - defusedxml

requirements_dev.txt

@@ -1,3 +1,4 @@
+albumentations
 black
 bump2version==0.5.11
 Click>=7.0

setup.py

@@ -11,6 +11,7 @@
     history = history_file.read()
 
 requirements = [
+    "albumentations>0.5.0",
     "Click>=7.0",
     "numpy",
     "pillow",

tests/test_stainaugment.py

@@ -0,0 +1,83 @@
+"""Tests for stain augmentation code."""
+
+import pathlib
+
+import albumentations as alb
+import numpy as np
+import pytest
+
+from tiatoolbox.data import stainnorm_target
+from tiatoolbox.tools.stainaugment import StainAugmentation
+from tiatoolbox.tools.stainnorm import get_normaliser
+from tiatoolbox.utils.misc import imread
+
+
+def test_stainaugment(source_image, norm_vahadane):
+    """Test functionality of the StainAugmentation class."""
+    source_img = imread(pathlib.Path(source_image))
+    target_img = stainnorm_target()
+    vahadane_img = imread(pathlib.Path(norm_vahadane))
+
+    # Test invalid method in the input
+    with pytest.raises(ValueError, match=r".*Unsupported stain extractor method.*"):
+        _ = StainAugmentation(method="mosi")
+
+    # 1. Testing without stain matrix.
+    # Test with Vahadane stain extractor
+    augmentor = StainAugmentation(
+        method="macenko", sigma1=3.0, sigma2=3.0, augment_background=True
+    )
+    augmentor.fit(source_img)
+    source_img_aug = augmentor.augment()
+    assert source_img_aug.dtype == source_img.dtype
+    assert np.shape(source_img_aug) == np.shape(source_img)
+    assert np.mean(np.absolute(source_img_aug / 255.0 - source_img / 255.0)) > 1e-2
+
+    # 2. Testing with predefined stain matrix
+    # We first extract the stain matrix of the target image and try to augment the
+    # source image with respect to that image.
+    norm = get_normaliser("vahadane")
+    norm.fit(target_img)
+    target_stain_matrix = norm.stain_matrix_target
+
+    # Now we augment the the source image with sigma1=0, sigma2=0 to force the augmentor
+    # to act like a normalizer
+    augmentor = StainAugmentation(
+        method="vahadane",
+        stain_matrix=target_stain_matrix,
+        sigma1=0.0,
+        sigma2=0.0,
+        augment_background=False,
+    )
+    augmentor.fit(source_img, threshold=0.8)
+    source_img_aug = augmentor.augment()
+    assert np.mean(np.absolute(vahadane_img / 255.0 - source_img_aug / 255.0)) < 1e-1
+
+    # 3. Test in albumentation framework
+    # Using the same trick as before, augment the image with pre-defined stain matrix
+    # and sigma1,2 equal to 0. The output should be equal to stain normalized image.
+    aug_pipeline = alb.Compose(
+        [
+            StainAugmentation(
+                method="vahadane",
+                stain_matrix=target_stain_matrix,
+                sigma1=0.0,
+                sigma2=0.0,
+                always_apply=True,
+            )
+        ],
+        p=1,
+    )
+    source_img_aug = aug_pipeline(image=source_img)["image"]
+    assert np.mean(np.absolute(vahadane_img / 255.0 - source_img_aug / 255.0)) < 1e-1
+
+    # Test for albumentation helper functions
+    params = augmentor.get_transform_init_args_names()
+    augmentor.get_params_dependent_on_targets(params)
+    assert params == (
+        "method",
+        "stain_matrix",
+        "sigma1",
+        "sigma2",
+        "augment_background",
+    )

tiatoolbox/tools/__init__.py

@@ -19,4 +19,10 @@
 # ***** END GPL LICENSE BLOCK *****
 
 """This package contains various tools for working with WSIs."""
-from tiatoolbox.tools import patchextraction, stainextract, stainnorm
+from tiatoolbox.tools import (
+    patchextraction,
+    stainaugment,
+    stainextract,
+    stainnorm,
+    tissuemask,
+)

tiatoolbox/tools/stainaugment.py

@@ -0,0 +1,214 @@
+# ***** BEGIN GPL LICENSE BLOCK *****
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 2
+# of the License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software Foundation,
+# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+#
+# This file contains code inspired by StainTools
+# [https://github.com/Peter554/StainTools] written by Peter Byfield.
+#
+# The Original Code is Copyright (C) 2021, TIACentre, University of Warwick
+# All rights reserved.
+# ***** END GPL LICENSE BLOCK *****
+
+"""Staing augmentation"""
+import copy
+import random
+
+import numpy as np
+from albumentations.core.transforms_interface import ImageOnlyTransform
+
+from tiatoolbox.tools.stainnorm import get_normaliser
+from tiatoolbox.utils.misc import get_luminosity_tissue_mask
+
+
+class StainAugmentation(ImageOnlyTransform):
+    """Stain augmentation using predefined stain matrix or stain extraction methods.
+
+    This stain augmentation class can be used in 'albumentations' augmentation pipelines
+    as well as stand alone. There is an option to use predefined `stain_matrix` in the
+    input which enables the `StainAugmentation` to generate augmented images faster or
+    do stain normalization to a specific target `stain_matrix`. Having stain matrix
+    beforhand, we don't need to do dictionary learning for stain matrix extraction,
+    hence,speed up the stain augmentation/normalization process which makes it more
+    appropriate for one-the-fly stain augmentation/normalization.
+
+    Args:
+        method (str): The method to use for stain matrix and stain concentration
+            extraction. Can be either "vahadane" (default) or "macenko".
+        stain_matrix (:class:`numpy.ndarray`): Pre-extracted stain matrix of a target
+            image. This can be used for both on-the-fly stain normalization and faster
+            stain augmentation. User can use tools in `tiatoolbox.tools.stainextract`
+            to extract this information. If None (default), the stain matrix will be
+            automatically extracted using the method specified by user.
+        sigma1 (float): Controls the extent of the stain concentrations scale parameter
+            (`alpha` belonging to [1-sigma1, 1+sigma1] range). Default is 0.5.
+        sigma2 (float): Controls the extent of the stain concentrations shift parameter
+            (`beta` belonging to [-sigma2, sigma2] range). Default is 0.25.
+        augment_background (bool): Specifies whether to apply stain augmentation on the
+            background or not. Default is False, which indicates that only tissue region
+            will be stain augmented.
+        always_apply (False): For use with 'albumentations' pipeline. Please refer to
+            albumentations documentations for more information.
+        p (0.5): For use with 'albumentations' pipeline which specifies the probability
+            of using the augmentation in a 'albumentations' pipeline. . Please refer to
+            albumentations documentations for more information.
+    Attributes:
+        stain_normaliser: Fitted stain normalization class.
+        stain_matrix (:class:`numpy.ndarray`): extracted stain matrix from the image
+        source_concentrations (:class:`numpy.ndarray`): extracted stain
+            concentrations from the input image.
+        n_stains (int): number of stain channels in the stain concentrations.
+            Expected to be 2 for H&E stained images.
+        tissue_mask (:class:`numpy.ndarray`): tissue region mask in the image.
+    Examples:
+        >>> '''Using the stain augmentor in the 'albumentations' pipeline'''
+        >>> from tiatoolbox.tools.stainaugment import StainAugmentaiton
+        >>> import albumentations as A
+        >>> # Defining an examplar stain matrix as refrence
+        >>> stain_matrix = np.array([[0.91633014, -0.20408072, -0.34451435],
+        ...                [0.17669817, 0.92528011, 0.33561059]])
+        >>> # Define albumentations pipeline
+        >>> aug_pipline = A.Compose([
+        ...                         A.RandomRotate90(),
+        ...                         A.Flip(),
+        ...                         StainAugmentaiton(stain_matrix=stain_matrix)
+        ...                         ])
+        >>> # apply the albumentations pipeline on an image (RGB numpy unit8 type)
+        >>> img_aug = aug(image=img)['image']
+
+        >>> '''Using the stain augmentor stand alone'''
+        >>> from tiatoolbox.tools.stainaugment import StainAugmentaiton
+        >>> # Defining an examplar stain matrix as refrence
+        >>> stain_matrix = np.array([[0.91633014, -0.20408072, -0.34451435],
+        ...                [0.17669817, 0.92528011, 0.33561059]])
+        >>> # Instantiate the stain augmentor and fit it on an image
+        >>> stain_augmentor = StainAugmentation(stain_matrix=stain_matrix)
+        >>> stain_augmentor.fit(img)
+        >>> # Now using the fitted `stain_augmentor` in a loop to generate
+        >>> # several augmented instances from the same image.
+        >>> for i in range(10):
+        ...     img_aug = stain_augmentor.augment()
+    """
+
+    def __init__(
+        self,
+        method: str = "vahadane",
+        stain_matrix: np.ndarray = None,
+        sigma1: float = 0.5,
+        sigma2: float = 0.25,
+        augment_background: bool = False,
+        always_apply=False,
+        p=0.5,
+    ) -> np.ndarray:
+        super().__init__(always_apply=always_apply, p=p)
+
+        self.augment_background = augment_background
+        self.sigma1 = sigma1
+        self.sigma2 = sigma2
+        self.method = method
+        self.stain_matrix = stain_matrix
+
+        if self.method.lower() not in {"macenko", "vahadane"}:
+            raise ValueError(
+                f"Unsupported stain extractor method '{self.method}' for "
+                "StainAugmentation. Choose either 'vahadane' or 'macenko'."
+            )
+        self.stain_normaliser = get_normaliser(self.method.lower())
+
+        self.alpha = None
+        self.beta = None
+        self.img_shape = None
+        self.tissue_mask = None
+        self.n_stains = None
+        self.source_concentrations = None
+
+    def fit(self, img, threshold=0.85):
+        """Fit function to extract information needed for stain augmentation.
+
+        The `fit` function uses either 'Macenko' or 'Vahadane' stain extraction methods
+        to extract stain matrix and stain concentrations of the input image to be used
+        in the `augment` function.
+
+        Args:
+            img (:class:`numpy.ndarray`): RGB image in the form of uint8 numpy array.
+            threshold (float): The threshold value used to find tissue mask from the
+                luminosity component of the image. The found `tissue_mask` will be used
+                to filter out background area in stain augmentation process upon user
+                setting `augment_background=False`.
+        """
+        if self.stain_matrix is None:
+            self.stain_normaliser.fit(img)
+            self.stain_matrix = self.stain_normaliser.stain_matrix_target
+            self.source_concentrations = self.stain_normaliser.target_concentrations
+        else:
+            self.source_concentrations = self.stain_normaliser.get_concentrations(
+                img, self.stain_matrix
+            )
+        self.n_stains = self.source_concentrations.shape[1]
+        self.tissue_mask = get_luminosity_tissue_mask(img, threshold=threshold).ravel()
+        self.img_shape = img.shape
+
+    def augment(self):
+        """Return an augmented instance based on source stain concentrations.
+
+        Stain concentrations of the source image are altered (scaled and shifted)
+        based on the random alpha and beta paramters, and then an augmented image is
+        reconstructed from the altered concentrations.
+        All parameters needed for this part are calculated when calling fit() function.
+
+        Returns:
+            img_augmented (:class:`numpy.ndarray`): stain augmented image.
+        """
+        self.get_params()
+        augmented_concentrations = copy.deepcopy(self.source_concentrations)
+        for i in range(self.n_stains):
+            if self.augment_background:
+                augmented_concentrations[:, i] *= self.alpha
+                augmented_concentrations[:, i] += self.beta
+            else:
+                augmented_concentrations[self.tissue_mask, i] *= self.alpha
+                augmented_concentrations[self.tissue_mask, i] += self.beta
+        img_augmented = 255 * np.exp(
+            -1 * np.dot(augmented_concentrations, self.stain_matrix)
+        )
+        img_augmented = img_augmented.reshape(self.img_shape)
+        img_augmented = np.clip(img_augmented, 0, 255)
+        return np.uint8(img_augmented)
+
+    def apply(self, img, **params):  # alpha=None, beta=None,
+        """Call the `fit` and `augment` functions to generate an stain augmented image.
+
+        Args:
+            img (:class:`numpy.ndarray`): Input RGB image in the form of unit8 numpy
+                array.
+        Returns:
+            :class:`numpy.ndarray`: Stain augmented image with the same
+                size and format as the input img.
+        """
+        self.fit(img, threshold=0.85)
+        return self.augment()
+
+    def get_params(self):
+        """Returns randomly generated parameters based on input arguments."""
+        self.alpha = random.uniform(1 - self.sigma1, 1 + self.sigma1)
+        self.beta = random.uniform(-self.sigma2, self.sigma2)
+        return {}
+
+    def get_params_dependent_on_targets(self, params):
+        """Does nothing, added to resolve flake 8 error"""
+        return {}
+
+    def get_transform_init_args_names(self):
+        """Return the argument names for albumentations use."""
+        return ("method", "stain_matrix", "sigma1", "sigma2", "augment_background")