Merge pull request #25 from andreped/reinhard

carloalbertobarbano · web-flow · commit 8b02b9f84647 · 2022-12-02T16:54:28.000+01:00
Add reinhard color normalization support
diff --git a/README.md b/README.md
@@ -5,12 +5,11 @@
 [![Pip Downloads](https://img.shields.io/pypi/dm/torchstain?label=pip%20downloads&logo=python)](https://pypi.org/project/torchstain/)
 [![DOI](https://zenodo.org/badge/323590093.svg)](https://zenodo.org/badge/latestdoi/323590093)
 
-
-
 GPU-accelerated stain normalization tools for histopathological images. Compatible with PyTorch, TensorFlow, and Numpy.
 Normalization algorithms currently implemented:
 
 - Macenko et al. [\[1\]](#reference) (ported from [numpy implementation](https://github.com/schaugf/HEnorm_python))
+- Reinhard et al. [\[2\]](#reference) (only numpy & TensorFlow backend support)
 
 ## Installation
 
@@ -43,20 +42,19 @@ t_to_transform = T(to_transform)
 norm, H, E = normalizer.normalize(I=t_to_transform, stains=True)
 ```
 
-![alt text](result.png)
+![alt text](data/result.png)
 
 ## Implemented algorithms
 
 | Algorithm | numpy | torch | tensorflow |
 |-|-|-|-|
 | Macenko | &check; | &check; | &check; |
-
+| Reinhard | &check; | &cross; | &check; |
 
 ## Backend comparison
 
 Results with 10 runs per size on a Intel(R) Core(TM) i5-8365U CPU @ 1.60GHz
 
-
 |   size | numpy avg. time   | torch avg. time   | tf avg. time     |
 |--------|-------------------|-------------------|------------------|
 |    224 | 0.0182s ± 0.0016  | 0.0180s ± 0.0390  | 0.0048s ± 0.0002 |
@@ -68,17 +66,15 @@ Results with 10 runs per size on a Intel(R) Core(TM) i5-8365U CPU @ 1.60GHz
 |   1568 | 1.1935s ± 0.0739  | 0.2590s ± 0.0088  | 0.2531s ± 0.0031 |
 |   1792 | 1.4523s ± 0.0207  | 0.3402s ± 0.0114  | 0.3080s ± 0.0188 |
 
-
 ## Reference
 
 - [1] Macenko, Marc, et al. "A method for normalizing histology slides for quantitative analysis." 2009 IEEE International Symposium on Biomedical Imaging: From Nano to Macro. IEEE, 2009.
-
+- [2] Reinhard, Erik, et al. "Color transfer between images." IEEE Computer Graphics and Applications. IEEE, 2001.
 
 ## Citing
 
 If you find this software useful for your research, please cite it as: 
 
-
 ```bibtex
 @software{barbano2022torchstain,
   author       = {Carlo Alberto Barbano and
diff --git a/data/result.png b/data/result.png
diff --git a/tests/test_color_conv.py b/tests/test_color_conv.py
@@ -0,0 +1,15 @@
+from torchstain.numpy.utils.rgb2lab import rgb2lab
+from torchstain.numpy.utils.lab2rgb import lab2rgb
+import numpy as np
+import cv2
+import os
+
+def test_rgb_to_lab():
+    size = 1024
+    curr_file_path = os.path.dirname(os.path.realpath(__file__))
+    img = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/source.png")), cv2.COLOR_BGR2RGB), (size, size))
+
+    reconstructed_img = lab2rgb(rgb2lab(img))
+    val = np.mean(np.abs(reconstructed_img - img))
+    print("MAE:", val)
+    assert val < 0.1
diff --git a/tests/test_torch.py b/tests/test_torch.py
@@ -9,7 +9,6 @@
 from torchvision import transforms
 from skimage.metrics import structural_similarity as ssim
 
-
 def setup_function(fn):
     print("torch version:", torch.__version__, "torchvision version:", torchvision.__version__)
 
diff --git a/torchstain/base/normalizers/__init__.py b/torchstain/base/normalizers/__init__.py
@@ -1,2 +1,3 @@
 from .he_normalizer import HENormalizer
 from .macenko import MacenkoNormalizer
+from .reinhard import ReinhardNormalizer
diff --git a/torchstain/base/normalizers/reinhard.py b/torchstain/base/normalizers/reinhard.py
@@ -0,0 +1,11 @@
+def ReinhardNormalizer(backend='numpy'):
+    if backend == 'numpy':
+        from torchstain.numpy.normalizers import NumpyReinhardNormalizer
+        return NumpyReinhardNormalizer()
+    elif backend == "torch":
+        raise NotImplementedError
+    elif backend == "tensorflow":
+        from torchstain.tf.normalizers import TensorFlowReinhardNormalizer
+        return TensorFlowReinhardNormalizer()
+    else:
+        raise Exception(f'Unknown backend {backend}')
diff --git a/torchstain/numpy/__init__.py b/torchstain/numpy/__init__.py
@@ -0,0 +1 @@
+from torchstain.numpy import normalizers, utils
diff --git a/torchstain/numpy/normalizers/__init__.py b/torchstain/numpy/normalizers/__init__.py
@@ -1 +1,2 @@
-from .macenko import NumpyMacenkoNormalizer
+from .macenko import NumpyMacenkoNormalizer
+from .reinhard import NumpyReinhardNormalizer
diff --git a/torchstain/numpy/normalizers/reinhard.py b/torchstain/numpy/normalizers/reinhard.py
@@ -0,0 +1,55 @@
+import numpy as np
+from torchstain.base.normalizers import HENormalizer
+from torchstain.numpy.utils.rgb2lab import rgb2lab
+from torchstain.numpy.utils.lab2rgb import lab2rgb
+from torchstain.numpy.utils.split import csplit, cmerge, lab_split, lab_merge
+from torchstain.numpy.utils.stats import get_mean_std, standardize
+
+"""
+Source code adapted from:
+https://github.com/DigitalSlideArchive/HistomicsTK/blob/master/histomicstk/preprocessing/color_normalization/reinhard.py
+https://github.com/Peter554/StainTools/blob/master/staintools/reinhard_color_normalizer.py
+"""
+class NumpyReinhardNormalizer(HENormalizer):
+    def __init__(self):
+        super().__init__()
+        self.target_mus = None
+        self.target_stds = None
+    
+    def fit(self, target):
+        # normalize
+        target = target.astype("float32") / 255
+
+        # convert to LAB
+        lab = rgb2lab(target)
+
+        # get summary statistics
+        stack_ = np.array([get_mean_std(x) for x in lab_split(lab)])
+        self.target_means = stack_[:, 0]
+        self.target_stds = stack_[:, 1]
+
+    def normalize(self, I):
+        # normalize
+        I = I.astype("float32") / 255
+        
+        # convert to LAB
+        lab = rgb2lab(I)
+        labs = lab_split(lab)
+
+        # get summary statistics from LAB
+        stack_ = np.array([get_mean_std(x) for x in labs])
+        mus = stack_[:, 0]
+        stds = stack_[:, 1]
+
+        # standardize intensities channel-wise and normalize using target mus and stds
+        result = [standardize(x, mu_, std_) * std_T + mu_T for x, mu_, std_, mu_T, std_T \
+            in zip(labs, mus, stds, self.target_means, self.target_stds)]
+        
+        # rebuild LAB
+        lab = lab_merge(*result)
+
+        # convert back to RGB from LAB
+        lab = lab2rgb(lab)
+
+        # rescale to [0, 255] uint8
+        return (lab * 255).astype("uint8")
diff --git a/torchstain/numpy/utils/__init__.py b/torchstain/numpy/utils/__init__.py
@@ -0,0 +1,4 @@
+from torchstain.numpy.utils.rgb2lab import *
+from torchstain.numpy.utils.lab2rgb import *
+from torchstain.numpy.utils.split import *
+from torchstain.numpy.utils.stats import *
diff --git a/torchstain/numpy/utils/lab2rgb.py b/torchstain/numpy/utils/lab2rgb.py
@@ -0,0 +1,38 @@
+import numpy as np
+from torchstain.numpy.utils.rgb2lab import _rgb2xyz
+
+_xyz2rgb = np.linalg.inv(_rgb2xyz)
+
+"""
+Implementation is based on:
+https://github.com/scikit-image/scikit-image/blob/00177e14097237ef20ed3141ed454bc81b308f82/skimage/color/colorconv.py#L704
+"""
+def lab2rgb(lab):
+    lab = lab.astype("float32")
+    # first rescale back from OpenCV format
+    lab[..., 0] /= 2.55
+    lab[..., 1] -= 128
+    lab[..., 2] -= 128
+
+    # convert LAB -> XYZ color domain
+    L, a, b = lab[..., 0], lab[..., 1], lab[..., 2]
+    y = (L + 16.) / 116.
+    x = (a / 500.) + y
+    z = y - (b / 200.)
+
+    out = np.stack([x, y, z], axis=-1)
+
+    mask = out > 0.2068966
+    out[mask] = np.power(out[mask], 3.)
+    out[~mask] = (out[~mask] - 16.0 / 116.) / 7.787
+
+    # rescale to the reference white (illuminant)
+    out *= np.array((0.95047, 1., 1.08883), dtype=out.dtype)
+    
+    # convert XYZ -> RGB color domain
+    arr = out.copy()
+    arr = np.dot(arr, _xyz2rgb.T)
+    mask = arr > 0.0031308
+    arr[mask] = 1.055 * np.power(arr[mask], 1 / 2.4) - 0.055
+    arr[~mask] *= 12.92
+    return np.clip(arr, 0, 1)
diff --git a/torchstain/numpy/utils/rgb2lab.py b/torchstain/numpy/utils/rgb2lab.py
@@ -0,0 +1,45 @@
+import numpy as np
+
+# constant conversion matrices between color spaces: https://gist.github.com/bikz05/6fd21c812ef6ebac66e1
+_rgb2xyz = np.array([[0.412453, 0.357580, 0.180423],
+                     [0.212671, 0.715160, 0.072169],
+                     [0.019334, 0.119193, 0.950227]])
+
+"""
+Implementation adapted from:
+https://gist.github.com/bikz05/6fd21c812ef6ebac66e1
+https://github.com/scikit-image/scikit-image/blob/00177e14097237ef20ed3141ed454bc81b308f82/skimage/color/colorconv.py#L704
+"""
+def rgb2lab(rgb):
+    rgb = rgb.astype("float32")
+
+    # convert rgb -> xyz color domain
+    arr = rgb.copy()
+    mask = arr > 0.04045
+    arr[mask] = np.power((arr[mask] + 0.055) / 1.055, 2.4)
+    arr[~mask] /= 12.92
+    xyz = np.dot(arr, _rgb2xyz.T.astype(arr.dtype))
+
+    # scale by CIE XYZ tristimulus values of the reference white point
+    arr = xyz.copy()
+    arr = arr / np.asarray((0.95047, 1., 1.08883), dtype=xyz.dtype)
+
+    # Nonlinear distortion and linear transformation
+    mask = arr > 0.008856
+    arr[mask] = np.cbrt(arr[mask])
+    arr[~mask] = 7.787 * arr[~mask] + 16. / 116.
+
+    x, y, z = arr[..., 0], arr[..., 1], arr[..., 2]
+
+    # Vector scaling
+    L = (116. * y) - 16.
+    a = 500.0 * (x - y)
+    b = 200.0 * (y - z)
+
+    # OpenCV format
+    L *= 2.55
+    a += 128
+    b += 128
+
+    # finally, get LAB color domain
+    return np.concatenate([x[..., np.newaxis] for x in [L, a, b]], axis=-1)
diff --git a/torchstain/numpy/utils/split.py b/torchstain/numpy/utils/split.py
@@ -0,0 +1,16 @@
+import numpy as np
+from torchstain.numpy.utils.rgb2lab import rgb2lab
+
+def csplit(I):
+    return [I[..., i] for i in range(I.shape[-1])]
+
+def cmerge(I1, I2, I3):
+    return np.stack([I1, I2, I3], axis=-1)
+
+def lab_split(I):
+    I = I.astype("float32")
+    I1, I2, I3 = csplit(I)
+    return I1 / 2.55, I2 - 128, I3 - 128
+
+def lab_merge(I1, I2, I3):
+    return cmerge(I1 * 2.55, I2 + 128, I3 + 128)
diff --git a/torchstain/numpy/utils/stats.py b/torchstain/numpy/utils/stats.py
@@ -0,0 +1,7 @@
+import numpy as np
+
+def get_mean_std(I):
+    return np.mean(I), np.std(I)
+
+def standardize(x, mu, std):
+    return (x - mu) / std
diff --git a/torchstain/tf/normalizers/__init__.py b/torchstain/tf/normalizers/__init__.py
@@ -1 +1,2 @@
-from torchstain.tf.normalizers.macenko import TensorFlowMacenkoNormalizer
+from torchstain.tf.normalizers.macenko import TensorFlowMacenkoNormalizer
+from torchstain.tf.normalizers.reinhard import TensorFlowReinhardNormalizer
diff --git a/torchstain/tf/normalizers/reinhard.py b/torchstain/tf/normalizers/reinhard.py
@@ -0,0 +1,55 @@
+import tensorflow as tf
+from torchstain.base.normalizers import HENormalizer
+from torchstain.tf.utils.rgb2lab import rgb2lab
+from torchstain.tf.utils.lab2rgb import lab2rgb
+from torchstain.tf.utils.split import csplit, cmerge, lab_split, lab_merge
+from torchstain.tf.utils.stats import get_mean_std, standardize
+
+"""
+Source code adapted from:
+https://github.com/DigitalSlideArchive/HistomicsTK/blob/master/histomicstk/preprocessing/color_normalization/reinhard.py
+https://github.com/Peter554/StainTools/blob/master/staintools/reinhard_color_normalizer.py
+"""
+class TensorFlowReinhardNormalizer(HENormalizer):
+    def __init__(self):
+        super().__init__()
+        self.target_mus = None
+        self.target_stds = None
+    
+    def fit(self, target):
+        # normalize
+        target = tf.cast(target, tf.float32) / 255
+
+        # convert to LAB
+        lab = rgb2lab(target)
+
+        # get summary statistics
+        stack_ = tf.convert_to_tensor([get_mean_std(x) for x in lab_split(lab)])
+        self.target_means = stack_[:, 0]
+        self.target_stds = stack_[:, 1]
+
+    def normalize(self, I):
+        # normalize
+        I = tf.cast(I, tf.float32) / 255
+        
+        # convert to LAB
+        lab = rgb2lab(I)
+        labs = lab_split(lab)
+
+        # get summary statistics from LAB
+        stack_ = tf.convert_to_tensor([get_mean_std(x) for x in labs])
+        mus = stack_[:, 0]
+        stds = stack_[:, 1]
+
+        # standardize intensities channel-wise and normalize using target mus and stds
+        result = [standardize(x, mu_, std_) * std_T + mu_T for x, mu_, std_, mu_T, std_T \
+            in zip(labs, mus, stds, self.target_means, self.target_stds)]
+        
+        # rebuild LAB
+        lab = lab_merge(*result)
+
+        # convert back to RGB from LAB
+        lab = lab2rgb(lab)
+
+        # rescale to [0, 255] uint8
+        return tf.cast(lab * 255, tf.uint8)
diff --git a/torchstain/tf/utils/__init__.py b/torchstain/tf/utils/__init__.py
@@ -1,3 +1,7 @@
 from torchstain.tf.utils.cov import cov
 from torchstain.tf.utils.percentile import percentile
 from torchstain.tf.utils.solveLS import solveLS
+from torchstain.tf.utils.stats import *
+from torchstain.tf.utils.split import *
+from torchstain.tf.utils.rgb2lab import *
+from torchstain.tf.utils.lab2rgb import *
diff --git a/torchstain/tf/utils/lab2rgb.py b/torchstain/tf/utils/lab2rgb.py
@@ -0,0 +1,36 @@
+import tensorflow as tf
+from torchstain.tf.utils.rgb2lab import _rgb2xyz, _white
+
+_xyz2rgb = tf.linalg.inv(_rgb2xyz)
+
+def lab2rgb(lab):
+    lab = tf.cast(lab, tf.float32)
+    
+    # rescale back from OpenCV format and extract LAB channel
+    L, a, b = lab[..., 0] / 2.55, lab[..., 1] - 128, lab[..., 2] - 128
+
+    # vector scaling to produce X, Y, Z
+    y = (L + 16.) / 116.
+    x = (a / 500.) + y
+    z = y - (b / 200.)
+
+    # merge back to get reconstructed XYZ color image
+    out = tf.stack([x, y, z], axis=-1)
+
+    # apply boolean transforms
+    mask = out > 0.2068966
+    not_mask = tf.math.logical_not(mask)
+    out = tf.tensor_scatter_nd_update(out, tf.where(mask), tf.pow(tf.boolean_mask(out, mask), 3))
+    out = tf.tensor_scatter_nd_update(out, tf.where(not_mask), (tf.boolean_mask(out, not_mask) - 16 / 116) / 7.787)
+
+    # rescale to the reference white (illuminant)
+    out = out * tf.cast(_white, out.dtype)
+    
+    # convert XYZ -> RGB color domain
+    arr = tf.identity(out)
+    arr = arr @ tf.transpose(_xyz2rgb)
+    mask = arr > 0.0031308
+    not_mask = tf.math.logical_not(mask)
+    arr = tf.tensor_scatter_nd_update(arr, tf.where(mask), 1.055 * tf.pow(tf.boolean_mask(arr, mask), 1 / 2.4) - 0.055)
+    arr = tf.tensor_scatter_nd_update(arr, tf.where(not_mask), tf.boolean_mask(out, not_mask) * 12.92)
+    return tf.clip_by_value(arr, 0, 1)
diff --git a/torchstain/tf/utils/rgb2lab.py b/torchstain/tf/utils/rgb2lab.py
diff --git a/torchstain/tf/utils/split.py b/torchstain/tf/utils/split.py
diff --git a/torchstain/tf/utils/stats.py b/torchstain/tf/utils/stats.py

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`from .he_normalizer import HENormalizer`
`2`	`2`	`from .macenko import MacenkoNormalizer`
	`3`	`+from .reinhard import ReinhardNormalizer`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from torchstain.numpy import normalizers, utils`
Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`		`-from .macenko import NumpyMacenkoNormalizer`
	`1`	`+from .macenko import NumpyMacenkoNormalizer`
	`2`	`+from .reinhard import NumpyReinhardNormalizer`