Merge pull request #45 from andreped/augmentation

Macenko stain augmentation

.github/workflows/tests_full.yml

@@ -39,7 +39,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ windows-2019, ubuntu-20.04, macos-11 ]
+        os: [ windows-2019, ubuntu-20.04, macos-12 ]
         python-version: [ 3.7, 3.8, 3.9 ]
         tf-version: [2.7.0, 2.8.0, 2.9.0]
 
@@ -51,7 +51,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Download artifact
-        uses: actions/download-artifact@master
+        uses: actions/download-artifact@v3
         with:
           name: "Python wheel"
 
@@ -70,7 +70,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ windows-2019, ubuntu-20.04, macos-11 ]
+        os: [ windows-2019, ubuntu-20.04, macos-12 ]
         python-version: [ 3.6, 3.7, 3.8, 3.9 ]
         pytorch-version: [1.8.0, 1.9.0, 1.10.0, 1.11.0, 1.12.0, 1.13.0]
         exclude:
@@ -89,7 +89,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Download artifact
-        uses: actions/download-artifact@master
+        uses: actions/download-artifact@v3
         with:
           name: "Python wheel"
 

.github/workflows/tests_quick.yml

@@ -43,7 +43,7 @@ jobs:
           python-version: 3.8
 
       - name: Download artifact
-        uses: actions/download-artifact@master
+        uses: actions/download-artifact@v3
         with:
           name: "Python wheel"
 
@@ -69,7 +69,7 @@ jobs:
           python-version: 3.8
 
       - name: Download artifact
-        uses: actions/download-artifact@master
+        uses: actions/download-artifact@v3
         with:
           name: "Python wheel"
 

apps/example_aug.py

@@ -0,0 +1,96 @@
+import cv2
+import matplotlib.pyplot as plt
+import torchstain
+import torch
+from torchvision import transforms
+import time
+import os
+
+
+size = 1024
+dir_path = os.path.dirname(os.path.abspath(__file__))
+target = cv2.resize(cv2.cvtColor(cv2.imread(dir_path + "/../data/target.png"), cv2.COLOR_BGR2RGB), (size, size))
+to_transform = cv2.resize(cv2.cvtColor(cv2.imread(dir_path + "/../data/source.png"), cv2.COLOR_BGR2RGB), (size, size))
+
+T = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Lambda(lambda x: x*255)
+])
+
+t_to_transform = T(to_transform)
+
+# setup augmentors for the different backends
+augmentor = torchstain.augmentors.MacenkoAugmentor(backend='numpy')
+augmentor.fit(to_transform)
+
+tf_augmentor = torchstain.augmentors.MacenkoAugmentor(backend='tensorflow')
+tf_augmentor.fit(t_to_transform)
+
+torch_augmentor = torchstain.augmentors.MacenkoAugmentor(backend='torch')
+torch_augmentor.fit(t_to_transform)
+
+
+print("NUMPY" + "-"*20)
+
+plt.figure()
+plt.suptitle('numpy augmentor')
+plt.subplot(4, 4, 1)
+plt.title('Original')
+plt.axis('off')
+plt.imshow(to_transform)
+
+for i in range(16):
+    # generate augmented sample
+    result = augmentor.augment()
+
+    plt.subplot(4, 4, i + 1)
+    if i == 1:
+        plt.title('Augmented ->')
+    plt.axis('off')
+    plt.imshow(result)
+
+plt.show()
+
+
+print("TensorFlow (TF)" + "-"*20)
+
+plt.figure()
+plt.suptitle('tf augmentor')
+plt.subplot(4, 4, 1)
+plt.title('Original')
+plt.axis('off')
+plt.imshow(to_transform)
+
+for i in range(16):
+    # generate augmented sample
+    result = tf_augmentor.augment()
+
+    plt.subplot(4, 4, i + 1)
+    if i == 1:
+        plt.title('Augmented ->')
+    plt.axis('off')
+    plt.imshow(result)
+
+plt.show()
+
+
+print("Torch" + "-"*20)
+
+plt.figure()
+plt.suptitle('torch augmentor')
+plt.subplot(4, 4, 1)
+plt.title('Original')
+plt.axis('off')
+plt.imshow(to_transform)
+
+for i in range(16):
+    # generate augmented sample
+    result = torch_augmentor.augment()
+
+    plt.subplot(4, 4, i + 1)
+    if i == 1:
+        plt.title('Augmented ->')
+    plt.axis('off')
+    plt.imshow(result)
+
+plt.show()

example.py → apps/example_norm.py

@@ -4,11 +4,13 @@
 import torch
 from torchvision import transforms
 import time
+import os
 
 
 size = 1024
-target = cv2.resize(cv2.cvtColor(cv2.imread("./data/target.png"), cv2.COLOR_BGR2RGB), (size, size))
-to_transform = cv2.resize(cv2.cvtColor(cv2.imread("./data/source.png"), cv2.COLOR_BGR2RGB), (size, size))
+dir_path = os.path.dirname(os.path.abspath(__file__))
+target = cv2.resize(cv2.cvtColor(cv2.imread(dir_path + "/../data/target.png"), cv2.COLOR_BGR2RGB), (size, size))
+to_transform = cv2.resize(cv2.cvtColor(cv2.imread(dir_path + "/../data/source.png"), cv2.COLOR_BGR2RGB), (size, size))
 
 normalizer = torchstain.normalizers.MacenkoNormalizer(backend='numpy')
 normalizer.fit(target)

torchstain/__init__.py

@@ -1,3 +1,3 @@
 __version__ = '1.3.0'
 
-from torchstain.base import normalizers
+from torchstain.base import augmentors, normalizers

torchstain/base/__init__.py

@@ -1 +1 @@
-from torchstain.base import normalizers
+from torchstain.base import augmentors, normalizers

torchstain/base/augmentors/__init__.py

@@ -0,0 +1,2 @@
+from .he_augmentor import HEAugmentor
+from .macenko import MacenkoAugmentor

torchstain/base/augmentors/he_augmentor.py

@@ -0,0 +1,6 @@
+class HEAugmentor:
+    def fit(self, I):
+        pass
+
+    def augment(self):
+        raise Exception('Abstract method')

torchstain/base/augmentors/macenko.py

@@ -0,0 +1,12 @@
+def MacenkoAugmentor(backend='torch', sigma1=0.2, sigma2=0.2):
+    if backend == 'numpy':
+        from torchstain.numpy.augmentors import NumpyMacenkoAugmentor
+        return NumpyMacenkoAugmentor(sigma1=sigma1, sigma2=sigma2)
+    elif backend == "torch":
+        from torchstain.torch.augmentors import TorchMacenkoAugmentor
+        return TorchMacenkoAugmentor(sigma1=sigma1, sigma2=sigma2)
+    elif backend == "tensorflow":
+        from torchstain.tf.augmentors import TensorFlowMacenkoAugmentor
+        return TensorFlowMacenkoAugmentor(sigma1=sigma1, sigma2=sigma2)
+    else:
+        raise Exception(f'Unknown backend {backend}')

torchstain/numpy/__init__.py

@@ -1 +1 @@
-from torchstain.numpy import normalizers, utils
+from torchstain.numpy import augmentors, normalizers, utils

torchstain/numpy/augmentors/__init__.py

@@ -0,0 +1 @@
+from .macenko import NumpyMacenkoAugmentor

torchstain/numpy/augmentors/macenko.py

@@ -0,0 +1,108 @@
+import numpy as np
+from torchstain.base.augmentors import HEAugmentor
+
+"""
+Source code adapted from: https://github.com/schaugf/HEnorm_python
+Original implementation: https://github.com/mitkovetta/staining-normalization
+"""
+class NumpyMacenkoAugmentor(HEAugmentor):
+    def __init__(self, sigma1=0.2, sigma2=0.2):
+        super().__init__()
+
+        self.sigma1 = sigma1
+        self.sigma2 = sigma2
+
+        self.I = None
+
+        self.HERef = np.array([[0.5626, 0.2159],
+                               [0.7201, 0.8012],
+                               [0.4062, 0.5581]])
+        self.maxCRef = np.array([1.9705, 1.0308])
+
+    def __convert_rgb2od(self, I, Io=240, beta=0.15):
+        # calculate optical density
+        OD = -np.log((I.astype(float) + 1) / Io)
+
+        # remove transparent pixels
+        ODhat = OD[~np.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 = ODhat.dot(eigvecs[:,1:3])
+
+        phi = np.arctan2(That[:,1],That[:,0])
+
+        minPhi = np.percentile(phi, alpha)
+        maxPhi = np.percentile(phi, 100-alpha)
+
+        vMin = eigvecs[:, 1:3].dot(np.array([(np.cos(minPhi), np.sin(minPhi))]).T)
+        vMax = eigvecs[:, 1:3].dot(np.array([(np.cos(maxPhi), np.sin(maxPhi))]).T)
+
+        # a heuristic to make the vector corresponding to hematoxylin first and the
+        # one corresponding to eosin second
+        if vMin[0] > vMax[0]:
+            HE = np.array((vMin[:,0], vMax[:,0])).T
+        else:
+            HE = np.array((vMax[:,0], vMin[:,0])).T
+
+        return HE
+
+    def __find_concentration(self, OD, HE):
+        # rows correspond to channels (RGB), columns to OD values
+        Y = np.reshape(OD, (-1, 3)).T
+
+        # determine concentrations of the individual stains
+        C = np.linalg.lstsq(HE, Y, rcond=None)[0]
+
+        return C
+
+    def __compute_matrices(self, I, Io, alpha, beta):
+        I = I.reshape((-1, 3))
+
+        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)
+
+        # compute eigenvectors
+        _, eigvecs = np.linalg.eigh(np.cov(ODhat.T))
+
+        HE = self.__find_HE(ODhat, eigvecs, alpha)
+
+        C = self.__find_concentration(OD, HE)
+
+        # normalize stain concentrations
+        maxC = np.array([np.percentile(C[0,:], 99), np.percentile(C[1,:], 99)])
+
+        return HE, C, maxC
+
+    def fit(self, I, Io=240, alpha=1, beta=0.15):
+        HE, C, maxC = self.__compute_matrices(I, Io, alpha, beta)
+
+        # keep these as we will use them for augmentation
+        self.I = I
+        self.HERef = HE
+        self.CRef = C
+        self.maxCRef = maxC
+
+    def augment(self, Io=240, alpha=1, beta=0.15):
+        I = self.I
+        h, w, c = I.shape
+        I = I.reshape((-1, 3))
+
+        HE, C, maxC = self.__compute_matrices(I, Io, alpha, beta)
+
+        maxC = np.divide(maxC, self.maxCRef)
+        C2 = np.divide(C, maxC[:, np.newaxis])
+
+        # introduce noise to the concentrations
+        for i in range(C2.shape[0]):
+            C2[i, :] *= np.random.uniform(1 - self.sigma1, 1 + self.sigma1)  # multiplicative
+            C2[i, :] += np.random.uniform(-self.sigma2, self.sigma2)  # additative
+
+        # recreate the image using reference mixing matrix
+        Iaug = np.multiply(Io, np.exp(-self.HERef.dot(C2)))
+        Iaug[Iaug > 255] = 255
+        Iaug = np.reshape(Iaug.T, (h, w, c)).astype(np.uint8)
+
+        return Iaug

torchstain/numpy/normalizers/macenko.py

@@ -10,8 +10,8 @@ def __init__(self):
         super().__init__()
 
         self.HERef = np.array([[0.5626, 0.2159],
-                          [0.7201, 0.8012],
-                          [0.4062, 0.5581]])
+                               [0.7201, 0.8012],
+                               [0.4062, 0.5581]])
         self.maxCRef = np.array([1.9705, 1.0308])
 
     def __convert_rgb2od(self, I, Io=240, beta=0.15):
@@ -109,7 +109,6 @@ def normalize(self, I, Io=240, alpha=1, beta=0.15, stains=True):
         Inorm[Inorm > 255] = 255
         Inorm = np.reshape(Inorm.T, (h, w, c)).astype(np.uint8)
 
-
         H, E = None, None
 
         if stains:

torchstain/tf/__init__.py

@@ -1 +1 @@
-from torchstain.tf import normalizers, utils
+from torchstain.tf import augmentors, normalizers, utils

torchstain/tf/augmentors/__init__.py

@@ -0,0 +1 @@
+from .macenko import TensorFlowMacenkoAugmentor