Refactor Stain Normalization

docs/source/apps.rst

@@ -100,15 +100,15 @@ Applications
     :members:
 
 .. automodule:: monai.apps.pathology.transforms.stain.array
-.. autoclass:: ExtractHEStains
+.. autoclass:: HEStainExtractor
     :members:
-.. autoclass:: NormalizeHEStains
+.. autoclass:: StainNormalizer
     :members:
 
 .. automodule:: monai.apps.pathology.transforms.stain.dictionary
-.. autoclass:: ExtractHEStainsd
+.. autoclass:: HEStainExtractord
     :members:
-.. autoclass:: NormalizeHEStainsd
+.. autoclass:: StainNormalizerd
     :members:
 
 .. automodule:: monai.apps.pathology.transforms.spatial.array

monai/apps/pathology/__init__.py

@@ -12,13 +12,13 @@
 from .data import MaskedInferenceWSIDataset, PatchWSIDataset, SmartCachePatchWSIDataset
 from .handlers import ProbMapProducer
 from .metrics import LesionFROC
-from .transforms.stain.array import ExtractHEStains, NormalizeHEStains
+from .transforms.stain.array import HEStainExtractor, StainNormalizer
 from .transforms.stain.dictionary import (
-    ExtractHEStainsd,
-    ExtractHEStainsD,
-    ExtractHEStainsDict,
-    NormalizeHEStainsd,
-    NormalizeHEStainsD,
-    NormalizeHEStainsDict,
+    HEStainExtractord,
+    HEStainExtractorD,
+    HEStainExtractorDict,
+    StainNormalizerd,
+    StainNormalizerD,
+    StainNormalizerDict,
 )
 from .utils import PathologyProbNMS, compute_isolated_tumor_cells, compute_multi_instance_mask

monai/apps/pathology/transforms/__init__.py

@@ -11,12 +11,12 @@
 
 from .spatial.array import SplitOnGrid, TileOnGrid
 from .spatial.dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict, TileOnGridd, TileOnGridD, TileOnGridDict
-from .stain.array import ExtractHEStains, NormalizeHEStains
+from .stain.array import HEStainExtractor, StainNormalizer
 from .stain.dictionary import (
-    ExtractHEStainsd,
-    ExtractHEStainsD,
-    ExtractHEStainsDict,
-    NormalizeHEStainsd,
-    NormalizeHEStainsD,
-    NormalizeHEStainsDict,
+    HEStainExtractord,
+    HEStainExtractorD,
+    HEStainExtractorDict,
+    StainNormalizerd,
+    StainNormalizerD,
+    StainNormalizerDict,
 )

monai/apps/pathology/transforms/stain/__init__.py

@@ -9,12 +9,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .array import ExtractHEStains, NormalizeHEStains
+from .array import HEStainExtractor, StainNormalizer
 from .dictionary import (
-    ExtractHEStainsd,
-    ExtractHEStainsD,
-    ExtractHEStainsDict,
-    NormalizeHEStainsd,
-    NormalizeHEStainsD,
-    NormalizeHEStainsDict,
+    HEStainExtractord,
+    HEStainExtractorD,
+    HEStainExtractorDict,
+    StainNormalizerd,
+    StainNormalizerD,
+    StainNormalizerDict,
 )

monai/apps/pathology/transforms/stain/array.py

@@ -16,181 +16,174 @@
 from monai.transforms.transform import Transform
 
 
-class ExtractHEStains(Transform):
-    """Class to extract a target stain from an image, using stain deconvolution (see Note).
+class HEStainExtractor(Transform):
+    """Extract stain coefficients from an image.
 
     Args:
-        tli: transmitted light intensity. Defaults to 240.
-        alpha: tolerance in percentile for the pseudo-min (alpha percentile)
-            and pseudo-max (100 - alpha percentile). Defaults to 1.
-        beta: absorbance threshold for transparent pixels. Defaults to 0.15
-        max_cref: reference maximum stain concentrations for Hematoxylin & Eosin (H&E).
-            Defaults to (1.9705, 1.0308).
+        source_intensity: transmitted light intensity.
+            Defaults to 240.
+        alpha: percentiles to ignore for outliers, so to calculate min and max,
+            if only consider (alpha, 100-alpha) percentiles. Defaults to 1.
+        beta: absorbance threshold for transparent pixels.
+            Defaults to 0.15
 
     Note:
-        For more information refer to:
-        - the original paper: Macenko et al., 2009 http://wwwx.cs.unc.edu/~mn/sites/default/files/macenko2009.pdf
-        - the previous implementations:
-
-          - MATLAB: https://github.com/mitkovetta/staining-normalization
-          - Python: https://github.com/schaugf/HEnorm_python
+        Please refer to this paper for further information on the method:
+        Macenko et al., 2009 http://wwwx.cs.unc.edu/~mn/sites/default/files/macenko2009.pdf
     """
 
-    def __init__(
-        self,
-        tli: float = 240,
-        alpha: float = 1,
-        beta: float = 0.15,
-        max_cref: Union[tuple, np.ndarray] = (1.9705, 1.0308),
-    ) -> None:
-        self.tli = tli
+    def __init__(self, source_intensity: float = 240, alpha: float = 1, beta: float = 0.15) -> None:
+        self.source_intensity = source_intensity
         self.alpha = alpha
         self.beta = beta
-        self.max_cref = np.array(max_cref)
 
-    def _deconvolution_extract_stain(self, image: np.ndarray) -> np.ndarray:
-        """Perform Stain Deconvolution and return stain matrix for the image.
+    def calculate_flat_absorbance(self, image):
+        """Calculate absorbace and remove transparent pixels"""
+        # calculate absorbance
+        image = image.astype(np.float32, copy=False) + 1.0
+        absorbance = -np.log(image.clip(max=self.source_intensity) / self.source_intensity)
+
+        # reshape to form a CxN matrix
+        c = absorbance.shape[0]
+        absorbance = absorbance.reshape((c, -1))
+
+        # remove transparent pixels
+        absorbance = absorbance[np.all(absorbance > self.beta, axis=1)]
+        if len(absorbance) == 0:
+            raise ValueError("All pixels of the input image are below the absorbance threshold.")
+
+        return absorbance
+
+    def _stain_decomposition(self, absorbance: np.ndarray) -> np.ndarray:
+        """Calculate the matrix of stain coefficient from the image.
 
         Args:
-            image: uint8 RGB image to perform stain deconvolution on
+            absorbance: absorbance matrix to perform stain extraction on
 
         Return:
-            he: H&E absorbance matrix for the image (first column is H, second column is E, rows are RGB values)
+            stain_coeff: stain attenuation coefficient matrix derive from the
+                image, where first column is H, second column is E, and
+                rows are RGB values
         """
-        # check image type and values
-        if not isinstance(image, np.ndarray):
-            raise TypeError("Image must be of type numpy.ndarray.")
-        if image.min() < 0:
-            raise ValueError("Image should not have negative values.")
-        if image.max() > 255:
-            raise ValueError("Image should not have values greater than 255.")
-
-        # reshape image and calculate absorbance
-        image = image.reshape((-1, 3))
-        image = image.astype(np.float32, copy=False) + 1.0
-        absorbance = -np.log(image.clip(max=self.tli) / self.tli)
-
-        # remove transparent pixels
-        absorbance_hat = absorbance[np.all(absorbance > self.beta, axis=1)]
-        if len(absorbance_hat) == 0:
-            raise ValueError("All pixels of the input image are below the absorbance threshold.")
 
         # compute eigenvectors
-        _, eigvecs = np.linalg.eigh(np.cov(absorbance_hat.T).astype(np.float32, copy=False))
+        _, eigvecs = np.linalg.eigh(np.cov(absorbance).astype(np.float32, copy=False))
 
         # project on the plane spanned by the eigenvectors corresponding to the two largest eigenvalues
-        t_hat = absorbance_hat.dot(eigvecs[:, 1:3])
+        projection = np.dot(eigvecs[:, -2:].T, absorbance)
 
-        # find the min and max vectors and project back to absorbance space
-        phi = np.arctan2(t_hat[:, 1], t_hat[:, 0])
+        # find the vectors that span the whole data (min and max angles)
+        phi = np.arctan2(projection[1], projection[0])
         min_phi = np.percentile(phi, self.alpha)
         max_phi = np.percentile(phi, 100 - self.alpha)
-        v_min = eigvecs[:, 1:3].dot(np.array([(np.cos(min_phi), np.sin(min_phi))], dtype=np.float32).T)
-        v_max = eigvecs[:, 1:3].dot(np.array([(np.cos(max_phi), np.sin(max_phi))], dtype=np.float32).T)
+        # project back to absorbance space
+        v_min = eigvecs[:, -2:].dot(np.array([(np.cos(min_phi), np.sin(min_phi))], dtype=np.float32).T)
+        v_max = eigvecs[:, -2:].dot(np.array([(np.cos(max_phi), np.sin(max_phi))], dtype=np.float32).T)
 
-        # a heuristic to make the vector corresponding to hematoxylin first and the one corresponding to eosin second
+        # make the vector corresponding to hematoxylin first and eosin second (based on R channel)
         if v_min[0] > v_max[0]:
-            he = np.array((v_min[:, 0], v_max[:, 0]), dtype=np.float32).T
+            stain_coeff = np.array((v_min[:, 0], v_max[:, 0]), dtype=np.float32).T
         else:
-            he = np.array((v_max[:, 0], v_min[:, 0]), dtype=np.float32).T
+            stain_coeff = np.array((v_max[:, 0], v_min[:, 0]), dtype=np.float32).T
 
-        return he
+        return stain_coeff
 
     def __call__(self, image: np.ndarray) -> np.ndarray:
         """Perform stain extraction.
 
         Args:
-            image: uint8 RGB image to extract stain from
+            image: RGB image to extract stain from
 
-        return:
-            target_he: H&E absorbance matrix for the image (first column is H, second column is E, rows are RGB values)
+        Return:
+            ref_stain_coeff: H&E absorbance matrix for the image (first column is H, second column is E, rows are RGB values)
         """
+        # check image type and values
         if not isinstance(image, np.ndarray):
-            raise TypeError("Image must be of type numpy.ndarray.")
+            raise TypeError("Image must be of type cupy.ndarray.")
+        if image.min() < 0:
+            raise ValueError("Image should not have negative values.")
 
-        target_he = self._deconvolution_extract_stain(image)
-        return target_he
+        absorbance = self.calculate_flat_absorbance(image)
+        ref_stain_coeff = self._stain_decomposition(absorbance)
+        return ref_stain_coeff
 
 
-class NormalizeHEStains(Transform):
-    """Class to normalize patches/images to a reference or target image stain (see Note).
+class StainNormalizer(Transform):
+    """Normalize images to a reference stain color matrix.
 
-    Performs stain deconvolution of the source image using the ExtractHEStains
-    class, to obtain the stain matrix and calculate the stain concentration matrix
-    for the image. Then, performs the inverse Beer-Lambert transform to recreate the
-    patch using the target H&E stain matrix provided. If no target stain provided, a default
-    reference stain is used. Similarly, if no maximum stain concentrations are provided, a
-    reference maximum stain concentrations matrix is used.
+    First, it extracts the stain coefficient matrix from the image using the provided stain extractor.
+    Then, it calculates the stain concentrations based on Beer-Lamber Law.
+    Next, it reconstructs the image using the provided reference stain matrix (stain-normalized image).
 
     Args:
-        tli: transmitted light intensity. Defaults to 240.
-        alpha: tolerance in percentile for the pseudo-min (alpha percentile) and
-            pseudo-max (100 - alpha percentile). Defaults to 1.
-        beta: absorbance threshold for transparent pixels. Defaults to 0.15.
-        target_he: target stain matrix. Defaults to ((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)).
-        max_cref: reference maximum stain concentrations for Hematoxylin & Eosin (H&E).
-            Defaults to [1.9705, 1.0308].
-
-    Note:
-        For more information refer to:
-        - the original paper: Macenko et al., 2009 http://wwwx.cs.unc.edu/~mn/sites/default/files/macenko2009.pdf
-        - the previous implementations:
+        source_intensity: transmitted light intensity.
+            Defaults to 240.
+        alpha: percentiles to ignore for outliers, so to calculate min and max,
+            if only consider (alpha, 100-alpha) percentiles. Defaults to 1.
+        ref_stain_coeff: reference stain attenuation coefficient matrix.
+            Defaults to ((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)).
+        ref_max_conc: reference maximum stain concentrations for
+            Hematoxylin & Eosin (H&E). Defaults to (1.9705, 1.0308).
 
-            - MATLAB: https://github.com/mitkovetta/staining-normalization
-            - Python: https://github.com/schaugf/HEnorm_python
     """
 
     def __init__(
         self,
-        tli: float = 240,
+        source_intensity: float = 240,
         alpha: float = 1,
-        beta: float = 0.15,
-        target_he: Union[tuple, np.ndarray] = ((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)),
-        max_cref: Union[tuple, np.ndarray] = (1.9705, 1.0308),
+        ref_stain_coeff: Union[tuple, np.ndarray] = ((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)),
+        ref_max_conc: Union[tuple, np.ndarray] = (1.9705, 1.0308),
+        stain_extractor=None,
     ) -> None:
-        self.tli = tli
-        self.target_he = np.array(target_he)
-        self.max_cref = np.array(max_cref)
-        self.stain_extractor = ExtractHEStains(tli=self.tli, alpha=alpha, beta=beta, max_cref=self.max_cref)
+        self.source_intensity = source_intensity
+        self.alpha = alpha
+        self.ref_stain_coeff = np.array(ref_stain_coeff)
+        self.ref_max_conc = np.array(ref_max_conc)
+        if stain_extractor is None:
+            self.stain_extractor = HEStainExtractor()
+        else:
+            self.stain_extractor = stain_extractor
 
     def __call__(self, image: np.ndarray) -> np.ndarray:
         """Perform stain normalization.
 
         Args:
-            image: uint8 RGB image/patch to be stain normalized, pixel values between 0 and 255
+            image: uint8 RGB image to be stain normalized, pixel values between 0 and 255
 
         Return:
             image_norm: stain normalized image/patch
         """
         # check image type and values
         if not isinstance(image, np.ndarray):
-            raise TypeError("Image must be of type numpy.ndarray.")
+            raise TypeError("Image must be of type cupy.ndarray.")
         if image.min() < 0:
             raise ValueError("Image should not have negative values.")
-        if image.max() > 255:
-            raise ValueError("Image should not have values greater than 255.")
 
-        # extract stain of the image
-        he = self.stain_extractor(image)
+        if self.source_intensity < 0:
+            raise ValueError("Source transmitted light intensity must be a positive value.")
+
+        # derive stain coefficient matrix from the image
+        stain_coeff = self.stain_extractor(image)
 
-        # reshape image and calculate absorbance
-        h, w, _ = image.shape
-        image = image.reshape((-1, 3))
-        image = image.astype(np.float32) + 1.0
-        absorbance = -np.log(image.clip(max=self.tli) / self.tli)
+        # calculate absorbance
+        image = image.astype(np.float32, copy=False) + 1.0
+        absorbance = -np.log(image.clip(max=self.source_intensity) / self.source_intensity)
 
-        # rows correspond to channels (RGB), columns to absorbance values
-        y = np.reshape(absorbance, (-1, 3)).T
+        # reshape to form a CxN matrix
+        c, h, w = absorbance.shape
+        absorbance = absorbance.reshape((c, -1))
 
-        # determine concentrations of the individual stains
-        conc = np.linalg.lstsq(he, y, rcond=None)[0]
+        # calculate concentrations of the each stain, based on Beer-Lambert Law
+        conc_raw = np.linalg.lstsq(stain_coeff, absorbance, rcond=None)[0]
 
         # normalize stain concentrations
-        max_conc = np.asarray([np.percentile(conc[0, :], 99), np.percentile(conc[1, :], 99)], dtype=np.float32)
-        tmp = np.divide(max_conc, self.max_cref, dtype=np.float32)
-        image_c = np.divide(conc, tmp[:, np.newaxis], dtype=np.float32)
-
-        image_norm: np.ndarray = np.multiply(self.tli, np.exp(-self.target_he.dot(image_c)), dtype=np.float32)
-        image_norm[image_norm > 255] = 254
-        image_norm = np.reshape(image_norm.T, (h, w, 3)).astype(np.uint8)
+        max_conc = np.percentile(conc_raw, 100 - self.alpha, axis=1)
+        normalization_factors = self.ref_max_conc / max_conc
+        conc_norm = conc_raw * normalization_factors[:, np.newaxis]
+
+        # reconstruct the image based on the reference stain matrix
+        image_norm: np.ndarray = np.multiply(
+            self.source_intensity, np.exp(-self.ref_stain_coeff.dot(conc_norm)), dtype=np.float32
+        )
+        image_norm = np.reshape(image_norm, (c, h, w)).astype(np.uint8)
         return image_norm