Merge pull request #5 from andreped/tensorflow-backend

Tensorflow backend support

Author: carloalbertobarbano

example.py

@@ -3,7 +3,8 @@
 import torchstain
 import torch
 from torchvision import transforms
-import numpy as np
+import time
+
 
 size = 1024
 target = cv2.resize(cv2.cvtColor(cv2.imread("./data/target.png"), cv2.COLOR_BGR2RGB), (size, size))
@@ -12,7 +13,6 @@
 normalizer = torchstain.MacenkoNormalizer(backend='numpy')
 normalizer.fit(target)
 
-
 T = transforms.Compose([
     transforms.ToTensor(),
     transforms.Lambda(lambda x: x*255)
@@ -21,9 +21,15 @@
 torch_normalizer = torchstain.MacenkoNormalizer(backend='torch')
 torch_normalizer.fit(T(target))
 
+tf_normalizer = torchstain.MacenkoNormalizer(backend='tensorflow')
+tf_normalizer.fit(T(target))
+
 t_to_transform = T(to_transform)
 
+t_ = time.time()
 norm, H, E = normalizer.normalize(I=to_transform, stains=True)
+print("numpy runtime:", time.time() - t_)
+
 plt.figure()
 plt.suptitle('numpy normalizer')
 plt.subplot(2, 2, 1)
@@ -47,7 +53,10 @@
 plt.imshow(E)
 plt.show()
 
+t_ = time.time()
 norm, H, E = torch_normalizer.normalize(I=t_to_transform, stains=True)
+print("torch runtime:", time.time() - t_)
+
 plt.figure()
 plt.suptitle('torch normalizer')
 plt.subplot(2, 2, 1)
@@ -70,3 +79,30 @@
 plt.axis('off')
 plt.imshow(E)
 plt.show()
+
+t_ = time.time()
+norm, H, E = tf_normalizer.normalize(I=t_to_transform, stains=True)
+print("tf runtime:", time.time() - t_)
+
+plt.figure()
+plt.suptitle('tensorflow normalizer')
+plt.subplot(2, 2, 1)
+plt.title('Original')
+plt.axis('off')
+plt.imshow(to_transform)
+
+plt.subplot(2, 2, 2)
+plt.title('Normalized')
+plt.axis('off')
+plt.imshow(norm)
+
+plt.subplot(2, 2, 3)
+plt.title('H')
+plt.axis('off')
+plt.imshow(H)
+
+plt.subplot(2, 2, 4)
+plt.title('E')
+plt.axis('off')
+plt.imshow(E)
+plt.show()

setup.py

@@ -18,7 +18,8 @@
     zip_safe=False,
     install_requires=[
         'torch',
-        'numpy'
+        'numpy',
+        'tensorflow'
     ],
     python_requires='>=3.6'
 )

torchstain/normalizers/macenko_normalizer.py

@@ -1,10 +1,13 @@
 from .numpy_macenko_normalizer import NumpyMacenkoNormalizer
 from .torch_macenko_normalizer import TorchMacenkoNormalizer
+from .tensorflow_macenko_normalizer import TensorFlowMacenkoNormalizer
 
 def MacenkoNormalizer(backend='torch'):
-    if backend not in ['torch', 'numpy']:
-        raise Exception(f'Unkown backend {backend}')
-
     if backend == 'numpy':
         return NumpyMacenkoNormalizer()
-    return TorchMacenkoNormalizer()
+    elif backend == "tensorflow":
+    	return TensorFlowMacenkoNormalizer()
+    elif backend == "torch":
+    	return TorchMacenkoNormalizer()
+    else:
+    	raise Exception(f'Unknown backend {backend}')

torchstain/normalizers/numpy_macenko_normalizer.py

@@ -50,7 +50,7 @@ def __find_concentration(self, OD, HE):
         Y = np.reshape(OD, (-1, 3)).T
 
         # determine concentrations of the individual stains
-        C = np.linalg.lstsq(HE,Y, rcond=None)[0]
+        C = np.linalg.lstsq(HE, Y, rcond=None)[0]
 
         return C
 

torchstain/normalizers/tensorflow_macenko_normalizer.py

@@ -0,0 +1,122 @@
+import tensorflow as tf
+from torchstain.normalizers.he_normalizer import HENormalizer
+from torchstain.utils import cov, percentile, cov_tf, percentile_tf, solveLS
+import numpy as np
+import tensorflow.keras.backend as K
+
+
+"""
+Source code ported from: https://github.com/schaugf/HEnorm_python
+Original implementation: https://github.com/mitkovetta/staining-normalization
+"""
+class TensorFlowMacenkoNormalizer(HENormalizer):
+    def __init__(self):
+        super().__init__()
+
+        self.HERef = tf.constant([[0.5626, 0.2159],
+                                   [0.7201, 0.8012],
+                                   [0.4062, 0.5581]])
+        self.maxCRef = tf.constant([1.9705, 1.0308])
+
+    def __convert_rgb2od(self, I, Io, beta):
+        I = tf.transpose(I, [1, 2, 0])
+
+        # calculate optical density
+        OD = -tf.math.log((tf.cast(tf.reshape(I, [tf.math.reduce_prod(I.shape[:-1]), I.shape[-1]]), tf.float32) + 1) / Io)
+
+        # remove transparent pixels
+        ODhat = OD[~tf.math.reduce_any(OD < beta, axis=1)]
+
+        return OD, ODhat
+
+    def __find_HE(self, ODhat, eigvecs, alpha):
+        # project on the plane spanned by the eigenvectors corresponding to the two
+        # largest eigenvalues
+        That = tf.linalg.matmul(ODhat, eigvecs)
+        phi = tf.math.atan2(That[:, 1], That[:, 0])
+
+        minPhi = percentile_tf(phi, alpha)
+        maxPhi = percentile_tf(phi, 100 - alpha)
+
+        vMin = tf.matmul(eigvecs, tf.expand_dims(tf.stack((tf.math.cos(minPhi), tf.math.sin(minPhi))), axis=-1))
+        vMax = tf.matmul(eigvecs, tf.expand_dims(tf.stack((tf.math.cos(maxPhi), tf.math.sin(maxPhi))), axis=-1))
+
+        # a heuristic to make the vector corresponding to hematoxylin first and the
+        # one corresponding to eosin second
+        HE = tf.where(vMin[0] > vMax[0], tf.concat((vMin, vMax), axis=1), tf.concat((vMax, vMin), axis=1))
+
+        return HE
+
+    def __find_concentration(self, OD, HE):
+        # rows correspond to channels (RGB), columns to OD values
+        Y = tf.transpose(OD)
+
+        # determine concentrations of the individual stains
+        return solveLS(HE, Y)[:2]
+
+    def __compute_matrices(self, I, Io, alpha, beta):
+        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)
+
+        # compute eigenvectors
+        _, eigvecs = tf.linalg.eigh(cov_tf(tf.transpose(ODhat)))
+        eigvecs = eigvecs[:, 1:3]
+
+        HE = self.__find_HE(ODhat, eigvecs, alpha)
+
+        C = self.__find_concentration(OD, HE)
+        maxC = tf.stack([percentile_tf(C[0, :], 99), percentile_tf(C[1, :], 99)])
+
+        return HE, C, maxC
+
+    def fit(self, I, Io=240, alpha=1, beta=0.15):
+        HE, _, maxC = self.__compute_matrices(I, Io, alpha, beta)
+
+        self.HERef = HE
+        self.maxCRef = maxC
+
+    def normalize(self, I, Io=240, alpha=1, beta=0.15, stains=True):
+        ''' Normalize staining appearence of H&E stained images
+
+        Example use:
+            see test.py
+
+        Input:
+            I: RGB input image: tensor of shape [C, H, W] and type uint8
+            Io: (optional) transmitted light intensity
+            alpha: percentile
+            beta: transparency threshold
+            stains: if true, return also H & E components
+
+        Output:
+            Inorm: normalized image
+            H: hematoxylin image
+            E: eosin image
+
+        Reference:
+            A method for normalizing histology slides for quantitative analysis. M.
+            Macenko et al., ISBI 2009
+        '''
+        c, h, w = I.shape
+
+        HE, C, maxC = self.__compute_matrices(I, Io, alpha, beta)
+
+        # normalize stain concentrations
+        C *= tf.expand_dims((self.maxCRef / maxC), axis=-1)
+
+        # recreate the image using reference mixing matrix
+        Inorm = Io * tf.math.exp(-tf.linalg.matmul(self.HERef, C))
+        Inorm = tf.clip_by_value(Inorm, 0, 255)
+        Inorm = tf.cast(tf.reshape(tf.transpose(Inorm), shape=(h, w, c)), tf.int32)
+
+        H, E = None, None
+
+        if stains:
+            H = tf.math.multiply(Io, tf.math.exp(tf.linalg.matmul(-tf.expand_dims(self.HERef[:, 0], axis=-1), tf.expand_dims(C[0, :], axis=0))))
+            H = tf.clip_by_value(H, 0, 255)
+            H = tf.cast(tf.reshape(tf.transpose(H), shape=(h, w, c)), tf.int32)
+
+            E = tf.math.multiply(Io, tf.math.exp(tf.linalg.matmul(-tf.expand_dims(self.HERef[:, 1], axis=-1), tf.expand_dims(C[1, :], axis=0))))
+            E = tf.clip_by_value(E, 0, 255)
+            E = tf.cast(tf.reshape(tf.transpose(E), shape=(h, w, c)), tf.int32)
+
+        return Inorm, H, E

torchstain/normalizers/torch_macenko_normalizer.py

@@ -20,7 +20,7 @@ def __convert_rgb2od(self, I, Io, beta):
         I = I.permute(1, 2, 0)
 
         # calculate optical density
-        OD = -torch.log((I.reshape((-1, I.shape[-1])).float()+1)/Io)
+        OD = -torch.log((I.reshape((-1, I.shape[-1])).float() + 1)/Io)
 
         # remove transparent pixels
         ODhat = OD[~torch.any(OD < beta, dim=1)]
@@ -36,8 +36,8 @@ def __find_HE(self, ODhat, eigvecs, alpha):
         minPhi = percentile(phi, alpha)
         maxPhi = percentile(phi, 100 - alpha)
 
-        vMin = torch.matmul(eigvecs, torch.stack((torch.cos(minPhi), torch.sin(minPhi))).T).unsqueeze(1)
-        vMax = torch.matmul(eigvecs, torch.stack((torch.cos(maxPhi), torch.sin(maxPhi))).T).unsqueeze(1)
+        vMin = torch.matmul(eigvecs, torch.stack((torch.cos(minPhi), torch.sin(minPhi)))).unsqueeze(1)
+        vMax = torch.matmul(eigvecs, torch.stack((torch.cos(maxPhi), torch.sin(maxPhi)))).unsqueeze(1)
 
         # a heuristic to make the vector corresponding to hematoxylin first and the
         # one corresponding to eosin second

torchstain/utils/__init__.py

@@ -1,2 +1,3 @@
-from torchstain.utils.cov import cov
-from torchstain.utils.percentile import percentile
+from torchstain.utils.cov import cov, cov_tf
+from torchstain.utils.percentile import percentile, percentile_tf
+from torchstain.utils.solveLS import solveLS

torchstain/utils/cov.py

@@ -1,4 +1,5 @@
 import torch
+import tensorflow as tf
 
 def cov(x):
     """
@@ -7,3 +8,12 @@ def cov(x):
     E_x = x.mean(dim=1)
     x = x - E_x[:, None]
     return torch.mm(x, x.T) / (x.size(1) - 1)
+
+
+def cov_tf(x):
+    """
+    https://en.wikipedia.org/wiki/Covariance_matrix
+    """
+    E_x = tf.math.reduce_mean(x, axis=1)
+    x = x - E_x[:, None]
+    return tf.linalg.matmul(x, tf.transpose(x)) / (x.shape[1] - 1)

torchstain/utils/percentile.py

@@ -1,5 +1,6 @@
 import torch
 from typing import Union
+import tensorflow as tf
 
 """
 Author: https://gist.github.com/spezold/42a451682422beb42bc43ad0c0967a30
@@ -21,5 +22,21 @@ def percentile(t: torch.tensor, q: float) -> Union[int, float]:
     # indeed corresponds to k=1, not k=0! Use float(q) instead of q directly,
     # so that ``round()`` returns an integer, even if q is a np.float32.
     k = 1 + round(.01 * float(q) * (t.numel() - 1))
-    result = t.view(-1).kthvalue(k).values
-    return result
+    return t.view(-1).kthvalue(k).values
+
+
+def percentile_tf(t: tf.Tensor, q: float) -> Union[int, float]:
+    """
+    Return the ``q``-th percentile of the flattened input tensor's data.
+
+    CAUTION:
+     * Needs PyTorch >= 1.1.0, as ``torch.kthvalue()`` is used.
+     * Values are not interpolated, which corresponds to
+       ``numpy.percentile(..., interpolation="nearest")``.
+
+    :param t: Input tensor.
+    :param q: Percentile to compute, which must be between 0 and 100 inclusive.
+    :return: Resulting value (scalar).
+    """
+    k = 1 + tf.math.round(.01 * tf.cast(q, tf.float32) * (tf.cast(tf.size(t), tf.float32) - 1))
+    return tf.sort(tf.reshape(t, [-1]))[tf.cast(k, tf.int32)]

torchstain/utils/solveLS.py

@@ -0,0 +1,8 @@
+import tensorflow as tf
+
+
+def solveLS(A, B):
+    q_full, _ = tf.linalg.qr(A, full_matrices=True)
+    ret1 = tf.linalg.lstsq(A, B)
+    ret2 = tf.linalg.lstsq(q_full, B)
+    return tf.concat([ret1, ret2[ret1.shape[0]:]], axis=0)