Repository: EIDOSLAB/torchstain (language: python)

Issue:
fixed tf percentile test
There was an off-by-one error in the in the kth-value stuff for the percentile implementation in TF.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index cf1fa66..f831348 100644
--- a/README.md
+++ b/README.md
@@ -3,8 +3,11 @@
 [![License](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
 [![tests](https://github.com/EIDOSLAB/torchstain/workflows/tests/badge.svg)](https://github.com/EIDOSLAB/torchstain/actions)
 [![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)
 
-Pytorch-compatible normalization tools for histopathological images.
+
+
+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))
@@ -12,9 +15,11 @@ Normalization algorithms currently implemented:
 ## Installation
 
 ```bash
-pip3 install torchstain
+pip install torchstain
 ```
 
+To install a specific backend use either ```torchstain[torch]``` or ```torchstain[tf]```. The numpy backend is included by default in both.
+
 ## Example Usage
 
 ```python
@@ -31,7 +36,7 @@ T = transforms.Compose([
     transforms.Lambda(lambda x: x*255)
 ])
 
-torch_normalizer = torchstain.MacenkoNormalizer(backend='torch')
+torch_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='torch')
 torch_normalizer.fit(T(target))
 
 t_to_transform = T(to_transform)
@@ -40,24 +45,50 @@ norm, H, E = normalizer.normalize(I=t_to_transform, stains=True)
 
 ![alt text](result.png)
 
+## Implemented algorithms
+
+| Algorithm | numpy | torch | tensorflow |
+|-|-|-|-|
+| Macenko | &check; | &check; | &check; |
+
 
 ## Backend comparison
 
 Results with 10 runs per size on a Intel(R) Core(TM) i5-8365U CPU @ 1.60GHz
 
 
-|   size | numpy avg. time   | numpy tot. time   | torch avg. time   | torch tot. time   |
-|--------|-------------------|-------------------|-------------------|-------------------|
-|    224x224 | 0.0323s ± 0.0032  | 0.3231s           | 0.0234s ± 0.0384  | 0.2340s           |
-|    448x448 | 0.1228s ± 0.0042  | 1.2280s           | 0.0395s ± 0.0168  | 0.3954s           |
-|    672x672 | 0.2653s ± 0.0106  | 2.6534s           | 0.0753s ± 0.0157  | 0.7527s           |
-|    896x896 | 0.4940s ± 0.0208  | 4.9397s           | 0.1262s ± 0.0159  | 1.2622s           |
-|   1120x1120 | 0.6888s ± 0.0081  | 6.8883s           | 0.2002s ± 0.0141  | 2.0021s           |
-|   1344x1344 | 1.0145s ± 0.0089  | 10.1448s          | 0.2703s ± 0.0136  | 2.7026s           |
-|   1568x1568 | 1.2620s ± 0.0133  | 12.6200s          | 0.3680s ± 0.0128  | 3.6795s           |
-|   1792x1792 | 1.4289s ± 0.0128  | 14.2886s          | 0.5968s ± 0.0160  | 5.9676s           |
+|   size | numpy avg. time   | torch avg. time   | tf avg. time     |
+|--------|-------------------|-------------------|------------------|
+|    224 | 0.0182s ± 0.0016  | 0.0180s ± 0.0390  | 0.0048s ± 0.0002 |
+|    448 | 0.0880s ± 0.0224  | 0.0283s ± 0.0172  | 0.0210s ± 0.0025 |
+|    672 | 0.1810s ± 0.0139  | 0.0463s ± 0.0301  | 0.0354s ± 0.0018 |
+|    896 | 0.3013s ± 0.0377  | 0.0820s ± 0.0329  | 0.0713s ± 0.0008 |
+|   1120 | 0.4694s ± 0.0350  | 0.1321s ± 0.0237  | 0.1036s ± 0.0042 |
+|   1344 | 0.6640s ± 0.0553  | 0.1665s ± 0.0026  | 0.1663s ± 0.0021 |
+|   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.
+
+
+## Citing
+
+If you find this software useful for your research, please cite it as: 
+
+
+```bibtex
+@software{barbano2022torchstain,
+  author       = {Carlo Alberto Barbano and
+                  André Pedersen},
+  title        = {EIDOSLAB/torchstain: v1.2.0-rc1},
+  month        = aug,
+  year         = 2022,
+  publisher    = {Zenodo},
+  version      = {v1.2.0-rc1},
+  doi          = {10.5281/zenodo.6976410},
+  url          = {https://doi.org/10.5281/zenodo.6976410}
+}
+```
diff --git a/compare.py b/compare.py
index 0fbf4f8..cacab8c 100644
--- a/compare.py
+++ b/compare.py
@@ -14,24 +14,26 @@ def measure(size, N):
     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))
 
-    normalizer = torchstain.MacenkoNormalizer(backend='numpy')
+    normalizer = torchstain.normalizers.MacenkoNormalizer(backend='numpy')
     normalizer.fit(target)
 
-
     T = transforms.Compose([
         transforms.ToPILImage(),
         transforms.ToTensor(),
         transforms.Lambda(lambda x: x*255)
     ])
 
-    torch_normalizer = torchstain.MacenkoNormalizer(backend='torch')
+    torch_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='torch')
     torch_normalizer.fit(T(target))
 
+    tf_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='tensorflow')
+    tf_normalizer.fit(T(target))
+
     t_to_transform = T(to_transform)
 
     t_np = []
     start_np = time.perf_counter()
-    for i in range(N):
+    for _ in range(N):
         tic = time.perf_counter()
         _ = normalizer.normalize(to_transform)
         toc = time.perf_counter()
@@ -42,7 +44,7 @@ def measure(size, N):
 
     t_torch = []
     start_torch = time.perf_counter()
-    for i in range(N):
+    for _ in range(N):
         tic = time.perf_counter()
         _ = torch_normalizer.normalize(t_to_transform)
         toc = time.perf_counter()
@@ -50,18 +52,32 @@ def measure(size, N):
     end_torch = time.perf_counter()
     t_torch = np.array(t_torch)
 
+
+    t_tf = []
+    start_tf = time.perf_counter()
+    for _ in range(N):
+        tic = time.perf_counter()
+        _ = torch_normalizer.normalize(t_to_transform)
+        toc = time.perf_counter()
+        t_tf.append(toc-tic)
+    end_tf = time.perf_counter()
+    t_tf = np.array(t_tf)
+
     """
     print(f'Results of {N} runs:')
     print(f'numpy: {t_np.mean():.4f}s ± {t_np.std():.4f} (tot: {end_np-start_np:.4f}s)')
     print(f'torch: {t_torch.mean():.4f}s ± {t_torch.std():.4f} (tot: {end_torch-start_torch:.4f}s)')
     """
 
-    return t_np, end_np-start_np, t_torch, end_torch-start_torch
+    return t_np, end_np-start_np, t_torch, end_torch-start_torch, t_tf, end_tf-start_tf
 
 table = []
 for size in [224, 448, 672, 896, 1120, 1344, 1568, 1792]:
-    t_np, tot_np, t_torch, tot_torch = measure(size, N=10)
-    row = [size, f'{t_np.mean():.4f}s ± {t_np.std():.4f}', f'{tot_np:.4f}s', f'{t_torch.mean():.4f}s ± {t_torch.std():.4f}', f'{tot_torch:.4f}s']
+    t_np, tot_np, t_torch, tot_torch, t_tf, tot_tf = measure(size, N=10)
+    # row = [size, f'{t_np.mean():.4f}s ± {t_np.std():.4f}', f'{tot_np:.4f}s', f'{t_torch.mean():.4f}s ± {t_torch.std():.4f}', f'{tot_torch:.4f}s']
+    row = [size, f'{t_np.mean():.4f}s ± {t_np.std():.4f}', f'{t_torch.mean():.4f}s ± {t_torch.std():.4f}', f'{t_tf.mean():.4f}s ± {t_tf.std():.4f}']
     table.append(row)
 
-print(tabulate(table, headers=['size', 'numpy avg. time', 'numpy tot. time', 'torch avg. time', 'torch tot. time'], tablefmt='github'))
+# print(tabulate(table, headers=['size', 'numpy avg. time', 'numpy tot. time', 'torch avg. time', 'torch tot. time'], tablefmt='github'))
+print(tabulate(table, headers=['size', 'numpy avg. time', 'torch avg. time', 'tf avg. time'], tablefmt='github'))
+
diff --git a/example.py b/example.py
index 4a6ea84..6de0f36 100644
--- a/example.py
+++ b/example.py
@@ -10,7 +10,7 @@
 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))
 
-normalizer = torchstain.MacenkoNormalizer(backend='numpy')
+normalizer = torchstain.normalizers.MacenkoNormalizer(backend='numpy')
 normalizer.fit(target)
 
 T = transforms.Compose([
@@ -18,10 +18,10 @@
     transforms.Lambda(lambda x: x*255)
 ])
 
-torch_normalizer = torchstain.MacenkoNormalizer(backend='torch')
+torch_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='torch')
 torch_normalizer.fit(T(target))
 
-tf_normalizer = torchstain.MacenkoNormalizer(backend='tensorflow')
+tf_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='tensorflow')
 tf_normalizer.fit(T(target))
 
 t_to_transform = T(to_transform)
diff --git a/setup.py b/setup.py
index 6f44501..d27acf7 100644
--- a/setup.py
+++ b/setup.py
@@ -1,13 +1,13 @@
 import pathlib
-from setuptools import setup, find_packages
+from setuptools import setup, find_packages, find_namespace_packages
 
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
 
 setup(
     name='torchstain',
-    version='1.1.0',
-    description='Pytorch stain normalization utils',
+    version='1.2.0',
+    description='Stain normalization tools for histological analysis and computational pathology',
     long_description=README,
     long_description_content_type='text/markdown',
     url='https://github.com/EIDOSlab/torchstain',
@@ -17,10 +17,12 @@
     packages=find_packages(exclude=('tests')),
     zip_safe=False,
     install_requires=[
-        'torch',
-        'numpy',
-        'tensorflow'
+        'numpy'
     ],
+    extras_require={
+        "tf": ["tensorflow"],
+        "torch": ["torch"],
+    },
     classifiers=[
         'Development Status :: 4 - Beta',
         'Intended Audience :: Developers',
diff --git a/torchstain/__init__.py b/torchstain/__init__.py
index ac31d26..dab8618 100644
--- a/torchstain/__init__.py
+++ b/torchstain/__init__.py
@@ -1,3 +1,3 @@
-__version__ = '1.1.0'
+__version__ = '1.2.0'
 
-from torchstain.normalizers.macenko_normalizer import MacenkoNormalizer
+from torchstain.base import normalizers
\ No newline at end of file
diff --git a/torchstain/base/__init__.py b/torchstain/base/__init__.py
new file mode 100644
index 0000000..d66541f
--- /dev/null
+++ b/torchstain/base/__init__.py
@@ -0,0 +1,1 @@
+from torchstain.base import normalizers
\ No newline at end of file
diff --git a/torchstain/base/normalizers/__init__.py b/torchstain/base/normalizers/__init__.py
new file mode 100644
index 0000000..aceb595
--- /dev/null
+++ b/torchstain/base/normalizers/__init__.py
@@ -0,0 +1,2 @@
+from .he_normalizer import HENormalizer
+from .macenko import MacenkoNormalizer
diff --git a/torchstain/normalizers/he_normalizer.py b/torchstain/base/normalizers/he_normalizer.py
similarity index 100%
rename from torchstain/normalizers/he_normalizer.py
rename to torchstain/base/normalizers/he_normalizer.py
diff --git a/torchstain/base/normalizers/macenko.py b/torchstain/base/normalizers/macenko.py
new file mode 100644
index 0000000..2c0c988
--- /dev/null
+++ b/torchstain/base/normalizers/macenko.py
@@ -0,0 +1,12 @@
+def MacenkoNormalizer(backend='torch'):
+    if backend == 'numpy':
+        from torchstain.numpy.normalizers import NumpyMacenkoNormalizer
+        return NumpyMacenkoNormalizer()
+    elif backend == "torch":
+        from torchstain.torch.normalizers.macenko import TorchMacenkoNormalizer
+        return TorchMacenkoNormalizer()
+    elif backend == "tensorflow":
+        from torchstain.tf.normalizers.macenko import TensorFlowMacenkoNormalizer
+        return TensorFlowMacenkoNormalizer()
+    else:
+        raise Exception(f'Unknown backend {backend}')
diff --git a/torchstain/normalizers/__init__.py b/torchstain/normalizers/__init__.py
deleted file mode 100644
index 4b24415..0000000
--- a/torchstain/normalizers/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-import torchstain.normalizers.macenko_normalizer
diff --git a/torchstain/normalizers/macenko_normalizer.py b/torchstain/normalizers/macenko_normalizer.py
deleted file mode 100644
index 266b550..0000000
--- a/torchstain/normalizers/macenko_normalizer.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from .numpy_macenko_normalizer import NumpyMacenkoNormalizer
-from .torch_macenko_normalizer import TorchMacenkoNormalizer
-from .tensorflow_macenko_normalizer import TensorFlowMacenkoNormalizer
-
-def MacenkoNormalizer(backend='torch'):
-    if backend == 'numpy':
-        return NumpyMacenkoNormalizer()
-    elif backend == "tensorflow":
-    	return TensorFlowMacenkoNormalizer()
-    elif backend == "torch":
-    	return TorchMacenkoNormalizer()
-    else:
-    	raise Exception(f'Unknown backend {backend}')
diff --git a/torchstain/numpy/__init__.py b/torchstain/numpy/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/torchstain/numpy/normalizers/__init__.py b/torchstain/numpy/normalizers/__init__.py
new file mode 100644
index 0000000..2220c20
--- /dev/null
+++ b/torchstain/numpy/normalizers/__init__.py
@@ -0,0 +1,1 @@
+from .macenko import NumpyMacenkoNormalizer
\ No newline at end of file
diff --git a/torchstain/normalizers/numpy_macenko_normalizer.py b/torchstain/numpy/normalizers/macenko.py
similarity index 97%
rename from torchstain/normalizers/numpy_macenko_normalizer.py
rename to torchstain/numpy/normalizers/macenko.py
index f32d702..37dac55 100644
--- a/torchstain/normalizers/numpy_macenko_normalizer.py
+++ b/torchstain/numpy/normalizers/macenko.py
@@ -1,5 +1,5 @@
 import numpy as np
-from torchstain.normalizers.he_normalizer import HENormalizer
+from torchstain.base.normalizers import HENormalizer
 
 """
 Source code adapted from: https://github.com/schaugf/HEnorm_python
@@ -16,7 +16,7 @@ def __init__(self):
 
     def __convert_rgb2od(self, I, Io=240, beta=0.15):
         # calculate optical density
-        OD = -np.log((I.astype(float)+1)/Io)
+        OD = -np.log((I.astype(np.float)+1)/Io)
 
         # remove transparent pixels
         ODhat = OD[~np.any(OD < beta, axis=1)]
diff --git a/torchstain/tf/__init__.py b/torchstain/tf/__init__.py
new file mode 100644
index 0000000..edaf586
--- /dev/null
+++ b/torchstain/tf/__init__.py
@@ -0,0 +1,1 @@
+from torchstain.tf import normalizers, utils
diff --git a/torchstain/tf/normalizers/__init__.py b/torchstain/tf/normalizers/__init__.py
new file mode 100644
index 0000000..f835e4d
--- /dev/null
+++ b/torchstain/tf/normalizers/__init__.py
@@ -0,0 +1,1 @@
+from torchstain.tf.normalizers.macenko import TensorFlowMacenkoNormalizer
\ No newline at end of file
diff --git a/torchstain/normalizers/tensorflow_macenko_normalizer.py b/torchstain/tf/normalizers/macenko.py
similarity index 91%
rename from torchstain/normalizers/tensorflow_macenko_normalizer.py
rename to torchstain/tf/normalizers/macenko.py
index 90ea859..bf21666 100644
--- a/torchstain/normalizers/tensorflow_macenko_normalizer.py
+++ b/torchstain/tf/normalizers/macenko.py
@@ -1,6 +1,6 @@
 import tensorflow as tf
-from torchstain.normalizers.he_normalizer import HENormalizer
-from torchstain.utils import cov, percentile, cov_tf, percentile_tf, solveLS
+from torchstain.base.normalizers.he_normalizer import HENormalizer
+from torchstain.tf.utils import cov, percentile, solveLS
 import numpy as np
 import tensorflow.keras.backend as K
 
@@ -35,8 +35,8 @@ def __find_HE(self, ODhat, eigvecs, alpha):
         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)
+        minPhi = percentile(phi, alpha)
+        maxPhi = percentile(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))
@@ -58,13 +58,13 @@ 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 = tf.linalg.eigh(cov(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)])
+        maxC = tf.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])
 
         return HE, C, maxC
 
diff --git a/torchstain/tf/utils/__init__.py b/torchstain/tf/utils/__init__.py
new file mode 100644
index 0000000..253a3d0
--- /dev/null
+++ b/torchstain/tf/utils/__init__.py
@@ -0,0 +1,3 @@
+from torchstain.tf.utils.cov import cov
+from torchstain.tf.utils.percentile import percentile
+from torchstain.tf.utils.solveLS import solveLS
diff --git a/torchstain/utils/cov.py b/torchstain/tf/utils/cov.py
similarity index 55%
rename from torchstain/utils/cov.py
rename to torchstain/tf/utils/cov.py
index 16dec9f..3429d90 100644
--- a/torchstain/utils/cov.py
+++ b/torchstain/tf/utils/cov.py
@@ -1,16 +1,6 @@
-import torch
 import tensorflow as tf
 
 def cov(x):
-    """
-    https://en.wikipedia.org/wiki/Covariance_matrix
-    """
-    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
     """
diff --git a/torchstain/tf/utils/percentile.py b/torchstain/tf/utils/percentile.py
new file mode 100644
index 0000000..32f7c67
--- /dev/null
+++ b/torchstain/tf/utils/percentile.py
@@ -0,0 +1,17 @@
+from typing import Union
+import tensorflow as tf
+
+def percentile(t: tf.Tensor, q: float) -> Union[int, float]:
+    """
+    Return the ``q``-th percentile of the flattened input tensor's data.
+
+    CAUTION:
+     * 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.math.reduce_prod(tf.size(t)), tf.float32) - 1))
+    return tf.sort(tf.reshape(t, [-1]))[tf.cast(k - 1, tf.int32)]
diff --git a/torchstain/utils/solveLS.py b/torchstain/tf/utils/solveLS.py
similarity index 99%
rename from torchstain/utils/solveLS.py
rename to torchstain/tf/utils/solveLS.py
index c3daf79..ee074aa 100644
--- a/torchstain/utils/solveLS.py
+++ b/torchstain/tf/utils/solveLS.py
@@ -1,6 +1,5 @@
 import tensorflow as tf
 
-
 def solveLS(A, B):
     q_full, _ = tf.linalg.qr(A, full_matrices=True)
     ret1 = tf.linalg.lstsq(A, B)
diff --git a/torchstain/torch/__init__.py b/torchstain/torch/__init__.py
new file mode 100644
index 0000000..fef6329
--- /dev/null
+++ b/torchstain/torch/__init__.py
@@ -0,0 +1,1 @@
+from torchstain.torch import normalizers, utils
\ No newline at end of file
diff --git a/torchstain/torch/normalizers/__init__.py b/torchstain/torch/normalizers/__init__.py
new file mode 100644
index 0000000..febcd90
--- /dev/null
+++ b/torchstain/torch/normalizers/__init__.py
@@ -0,0 +1,1 @@
+from torchstain.torch.normalizers.macenko import TorchMacenkoNormalizer
diff --git a/torchstain/normalizers/torch_macenko_normalizer.py b/torchstain/torch/normalizers/macenko.py
similarity index 97%
rename from torchstain/normalizers/torch_macenko_normalizer.py
rename to torchstain/torch/normalizers/macenko.py
index 8a0e8bc..710b7e3 100644
--- a/torchstain/normalizers/torch_macenko_normalizer.py
+++ b/torchstain/torch/normalizers/macenko.py
@@ -1,6 +1,6 @@
 import torch
-from torchstain.normalizers.he_normalizer import HENormalizer
-from torchstain.utils import cov, percentile
+from torchstain.base.normalizers.he_normalizer import HENormalizer
+from torchstain.torch.utils import cov, percentile
 
 """
 Source code ported from: https://github.com/schaugf/HEnorm_python
diff --git a/torchstain/torch/utils/__init__.py b/torchstain/torch/utils/__init__.py
new file mode 100644
index 0000000..5e0de3e
--- /dev/null
+++ b/torchstain/torch/utils/__init__.py
@@ -0,0 +1,2 @@
+from torchstain.torch.utils.cov import cov
+from torchstain.torch.utils.percentile import percentile
diff --git a/torchstain/torch/utils/cov.py b/torchstain/torch/utils/cov.py
new file mode 100644
index 0000000..5a2d3c1
--- /dev/null
+++ b/torchstain/torch/utils/cov.py
@@ -0,0 +1,9 @@
+import torch
+
+def cov(x):
+    """
+    https://en.wikipedia.org/wiki/Covariance_matrix
+    """
+    E_x = x.mean(dim=1)
+    x = x - E_x[:, None]
+    return torch.mm(x, x.T) / (x.size(1) - 1)
diff --git a/torchstain/torch/utils/percentile.py b/torchstain/torch/utils/percentile.py
new file mode 100644
index 0000000..08c28ef
--- /dev/null
+++ b/torchstain/torch/utils/percentile.py
@@ -0,0 +1,24 @@
+import torch
+from typing import Union
+
+"""
+Author: https://gist.github.com/spezold/42a451682422beb42bc43ad0c0967a30
+"""
+def percentile(t: torch.Tensor, q: float) -> Union[int, float]:
+    """
+    Return the ``q``-th percentile of the flattenepip d 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).
+    """
+    # Note that ``kthvalue()`` works one-based, i.e. the first sorted value
+    # 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))
+    return t.view(-1).kthvalue(k).values
diff --git a/torchstain/utils/__init__.py b/torchstain/utils/__init__.py
deleted file mode 100644
index 2318156..0000000
--- a/torchstain/utils/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from torchstain.utils.cov import cov, cov_tf
-from torchstain.utils.percentile import percentile, percentile_tf
-from torchstain.utils.solveLS import solveLS
diff --git a/torchstain/utils/percentile.py b/torchstain/utils/percentile.py
deleted file mode 100644
index 5bb263b..0000000
--- a/torchstain/utils/percentile.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import torch
-from typing import Union
-import tensorflow as tf
-
-"""
-Author: https://gist.github.com/spezold/42a451682422beb42bc43ad0c0967a30
-"""
-def percentile(t: torch.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).
-    """
-    # Note that ``kthvalue()`` works one-based, i.e. the first sorted value
-    # 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))
-    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)]


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
deleted file mode 100644
index cf18d53..0000000
--- a/.github/workflows/tests.yml
+++ /dev/null
@@ -1,65 +0,0 @@
-name: tests
-
-on:
-  push:
-    branches:
-      - '*'
-  pull_request:
-    branches:
-      - '*'
-
-jobs:
-  build:
-    runs-on: ubuntu-18.04
-    steps:
-      - uses: actions/checkout@v1
-      - name: Set up Python 3.6
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.6
-
-      - name: Install dependencies
-        run: pip install wheel setuptools torch tensorflow numpy
-
-      - name: Build wheel
-        run: python setup.py bdist_wheel
-
-      - name: Upload Python wheel
-        uses: actions/upload-artifact@v2
-        with:
-          name: Python wheel
-          path: ${{github.workspace}}/dist/torchstain-*.whl
-          if-no-files-found: error
-
-  test:
-    needs: build
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [windows-2019, ubuntu-18.04, macos-11]
-        python-version: [3.6, 3.7, 3.8, 3.9]
-
-    steps:
-      - uses: actions/checkout@v1
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        run: pip install wheel setuptools torch tensorflow numpy
-
-      - name: Install test dependencies
-        run: pip install opencv-python matplotlib torchvision scikit-image pytest
-
-      - name: Download artifact
-        uses: actions/download-artifact@master
-        with:
-          name: "Python wheel"
-
-      - name: Install wheel
-        run: pip install --find-links=${{github.workspace}} torchstain
-
-      - name: Run tests
-        run: |
-          pytest -v tests
diff --git a/.github/workflows/tests_full.yml b/.github/workflows/tests_full.yml
new file mode 100644
index 0000000..ba105d9
--- /dev/null
+++ b/.github/workflows/tests_full.yml
@@ -0,0 +1,101 @@
+name: tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-18.04
+    steps:
+      - uses: actions/checkout@v1
+      - name: Set up Python 3.6
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.6
+
+      - name: Install dependencies
+        run: pip install wheel setuptools
+
+      - name: Build wheel
+        run: python setup.py bdist_wheel
+
+      - name: Upload Python wheel
+        uses: actions/upload-artifact@v2
+        with:
+          name: Python wheel
+          path: ${{github.workspace}}/dist/torchstain-*.whl
+          if-no-files-found: error
+
+  test-tf:
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ windows-2019, ubuntu-18.04, macos-11 ]
+        python-version: [ 3.7, 3.8, 3.9 ]
+        tf-version: [2.7.0, 2.8.0, 2.9.0]
+
+    steps:
+      - uses: actions/checkout@v1
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Download artifact
+        uses: actions/download-artifact@master
+        with:
+          name: "Python wheel"
+
+      - name: Install dependencies
+        run: |
+          pip install tensorflow==${{ matrix.tf-version }} protobuf==3.20.* opencv-python-headless scikit-image
+          pip install pytest
+
+      - name: Install wheel
+        run: pip install --find-links=${{github.workspace}} torchstain
+
+      - name: Run tests
+        run: pytest -vs tests/test_tf.py
+
+  test-torch:
+    needs: build
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ windows-2019, ubuntu-18.04, macos-11 ]
+        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]
+        exclude:
+          - python-version: 3.6
+            pytorch-version: 1.11.0
+          - python-version: 3.6
+            pytorch-version: 1.12.0
+
+    steps:
+      - uses: actions/checkout@v1
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Download artifact
+        uses: actions/download-artifact@master
+        with:
+          name: "Python wheel"
+
+      - name: Install dependencies
+        run: |
+          pip install torch==${{ matrix.pytorch-version }} torchvision opencv-python-headless scikit-image
+          pip install pytest
+
+      - name: Install wheel
+        run: pip install --find-links=${{github.workspace}} torchstain
+
+      - name: Run tests
+        run: pytest -vs tests/test_torch.py
diff --git a/.github/workflows/tests_quick.yml b/.github/workflows/tests_quick.yml
new file mode 100644
index 0000000..3e08828
--- /dev/null
+++ b/.github/workflows/tests_quick.yml
@@ -0,0 +1,86 @@
+name: tests
+
+on:
+  push:
+    branches-ignore:
+      - main
+  pull_request:
+    branches-ignore:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-18.04
+    steps:
+      - uses: actions/checkout@v1
+      - name: Set up Python 3.6
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.6
+
+      - name: Install dependencies
+        run: pip install wheel setuptools
+
+      - name: Build wheel
+        run: python setup.py bdist_wheel
+
+      - name: Upload Python wheel
+        uses: actions/upload-artifact@v2
+        with:
+          name: Python wheel
+          path: ${{github.workspace}}/dist/torchstain-*.whl
+          if-no-files-found: error
+
+  test-tf:
+    needs: build
+    runs-on: ubuntu-18.04
+
+    steps:
+      - uses: actions/checkout@v1
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Download artifact
+        uses: actions/download-artifact@master
+        with:
+          name: "Python wheel"
+
+      - name: Install dependencies
+        run: |
+          pip install tensorflow protobuf==3.20.* opencv-python-headless scikit-image
+          pip install pytest
+
+      - name: Install wheel
+        run: pip install --find-links=${{github.workspace}} torchstain
+
+      - name: Run tests
+        run: pytest -vs tests/test_tf.py
+
+  test-torch:
+    needs: build
+    runs-on: ubuntu-18.04
+
+    steps:
+      - uses: actions/checkout@v1
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Download artifact
+        uses: actions/download-artifact@master
+        with:
+          name: "Python wheel"
+
+      - name: Install dependencies
+        run: |
+          pip install torch torchvision opencv-python-headless scikit-image
+          pip install pytest
+
+      - name: Install wheel
+        run: pip install --find-links=${{github.workspace}} torchstain
+
+      - name: Run tests
+        run: pytest -vs tests/test_torch.py
diff --git a/tests/test_normalizers.py b/tests/test_normalizers.py
deleted file mode 100644
index e3da989..0000000
--- a/tests/test_normalizers.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import os
-import cv2
-import matplotlib.pyplot as plt
-import torchstain
-import torch
-from torchvision import transforms
-import time
-from skimage.metrics import structural_similarity as ssim
-import numpy as np
-
-
-size = 1024
-curr_file_path = os.path.dirname(os.path.realpath(__file__))
-target = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/target.png")), cv2.COLOR_BGR2RGB), (size, size))
-to_transform = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/source.png")), cv2.COLOR_BGR2RGB), (size, size))
-
-# setup preprocessing and preprocess image to be normalized
-T = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Lambda(lambda x: x*255)
-])
-t_to_transform = T(to_transform)
-
-# initialize normalizers for each backend and fit to target image
-normalizer = torchstain.MacenkoNormalizer(backend='numpy')
-normalizer.fit(target)
-
-torch_normalizer = torchstain.MacenkoNormalizer(backend='torch')
-torch_normalizer.fit(T(target))
-
-tf_normalizer = torchstain.MacenkoNormalizer(backend='tensorflow')
-tf_normalizer.fit(T(target))
-
-# transform
-result_numpy, _, _ = normalizer.normalize(I=to_transform, stains=True)
-result_torch, _, _ = torch_normalizer.normalize(I=t_to_transform, stains=True)
-result_tf, _, _ = tf_normalizer.normalize(I=t_to_transform, stains=True)
-
-# convert to numpy and set dtype
-result_numpy = result_numpy.astype("float32")
-result_torch = result_torch.numpy().astype("float32")
-result_tf = result_tf.numpy().astype("float32")
-
-# assess whether the normalized images are identical across backends
-np.testing.assert_almost_equal(ssim(result_numpy.flatten(), result_torch.flatten()), 1.0, decimal=4, verbose=True)
-np.testing.assert_almost_equal(ssim(result_numpy.flatten(), result_tf.flatten()), 1.0, decimal=4, verbose=True)
diff --git a/tests/test_tf.py b/tests/test_tf.py
new file mode 100644
index 0000000..b24757a
--- /dev/null
+++ b/tests/test_tf.py
@@ -0,0 +1,51 @@
+import os
+import cv2
+import torchstain
+import torchstain.tf
+import tensorflow as tf
+import time
+from skimage.metrics import structural_similarity as ssim
+import numpy as np
+
+def test_cov():
+    x = np.random.randn(10, 10)
+    cov_np = np.cov(x)
+    cov_t = torchstain.tf.utils.cov(x)
+
+    np.testing.assert_almost_equal(cov_np, cov_t.numpy())
+
+def test_percentile():
+    x = np.random.randn(10, 10)
+    p = 20
+    p_np = np.percentile(x, p, interpolation='nearest')
+    p_t = torchstain.tf.utils.percentile(x, p)
+
+    np.testing.assert_almost_equal(p_np, p_t)
+
+def test_normalize_tf():
+    size = 1024
+    curr_file_path = os.path.dirname(os.path.realpath(__file__))
+    target = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/target.png")), cv2.COLOR_BGR2RGB), (size, size))
+    to_transform = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/source.png")), cv2.COLOR_BGR2RGB), (size, size))
+
+    # setup preprocessing and preprocess image to be normalized
+    T = lambda x: tf.convert_to_tensor(np.moveaxis(x, -1, 0).astype("float32"))  #  * 255
+    t_to_transform = T(to_transform)
+
+    # initialize normalizers for each backend and fit to target image
+    normalizer = torchstain.normalizers.MacenkoNormalizer(backend='numpy')
+    normalizer.fit(target)
+
+    tf_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='tensorflow')
+    tf_normalizer.fit(T(target))
+
+    # transform
+    result_numpy, _, _ = normalizer.normalize(I=to_transform, stains=True)
+    result_tf, _, _ = tf_normalizer.normalize(I=t_to_transform, stains=True)
+
+    # convert to numpy and set dtype
+    result_numpy = result_numpy.astype("float32")
+    result_tf = result_tf.numpy().astype("float32")
+
+    # assess whether the normalized images are identical across backends
+    np.testing.assert_almost_equal(ssim(result_numpy.flatten(), result_tf.flatten()), 1.0, decimal=4, verbose=True)
diff --git a/tests/test_torch.py b/tests/test_torch.py
new file mode 100644
index 0000000..f4363e7
--- /dev/null
+++ b/tests/test_torch.py
@@ -0,0 +1,60 @@
+import os
+import cv2
+import torchstain
+import torchstain.torch
+import torch
+import torchvision
+import time
+import numpy as np
+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__)
+
+def test_cov():
+    x = np.random.randn(10, 10)
+    cov_np = np.cov(x)
+    cov_t = torchstain.torch.utils.cov(torch.tensor(x))
+
+    np.testing.assert_almost_equal(cov_np, cov_t.numpy())
+
+def test_percentile():
+    x = np.random.randn(10, 10)
+    p = 20
+    p_np = np.percentile(x, p, interpolation='nearest')
+    p_t = torchstain.torch.utils.percentile(torch.tensor(x), p)
+
+    np.testing.assert_almost_equal(p_np, p_t)
+
+def test_normalize_torch():
+    size = 1024
+    curr_file_path = os.path.dirname(os.path.realpath(__file__))
+    target = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/target.png")), cv2.COLOR_BGR2RGB), (size, size))
+    to_transform = cv2.resize(cv2.cvtColor(cv2.imread(os.path.join(curr_file_path, "../data/source.png")), cv2.COLOR_BGR2RGB), (size, size))
+
+    # setup preprocessing and preprocess image to be normalized
+    T = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Lambda(lambda x: x*255)
+    ])
+    t_to_transform = T(to_transform)
+
+    # initialize normalizers for each backend and fit to target image
+    normalizer = torchstain.normalizers.MacenkoNormalizer(backend='numpy')
+    normalizer.fit(target)
+
+    torch_normalizer = torchstain.normalizers.MacenkoNormalizer(backend='torch')
+    torch_normalizer.fit(T(target))
+
+    # transform
+    result_numpy, _, _ = normalizer.normalize(I=to_transform, stains=True)
+    result_torch, _, _ = torch_normalizer.normalize(I=t_to_transform, stains=True)
+
+    # convert to numpy and set dtype
+    result_numpy = result_numpy.astype("float32")
+    result_torch = result_torch.numpy().astype("float32")
+
+    # assess whether the normalized images are identical across backends
+    np.testing.assert_almost_equal(ssim(result_numpy.flatten(), result_torch.flatten()), 1.0, decimal=4, verbose=True)
diff --git a/tests/test_utils.py b/tests/test_utils.py
deleted file mode 100644
index 0cf3912..0000000
--- a/tests/test_utils.py
+++ /dev/null
@@ -1,19 +0,0 @@
-import torch
-import torchstain
-import numpy as np
-
-
-def test_cov():
-    x = np.random.randn(10, 10)
-    cov_np = np.cov(x)
-    cov_t = torchstain.utils.cov(torch.tensor(x))
-
-    np.testing.assert_almost_equal(cov_np, cov_t.numpy())
-
-def test_percentile():
-    x = np.random.randn(10, 10)
-    p = 20
-    p_np = np.percentile(x, p, method='nearest')
-    p_t = torchstain.utils.percentile(torch.tensor(x), p)
-
-    np.testing.assert_almost_equal(p_np, p_t)