update unittest signatures; remove one invalid test case (concentration output is never masked)

Author: CielAl

README.md

@@ -153,7 +153,7 @@ augmentor = AugmentorBuilder.build('vahadane',
                                    # if set to None means all pixels are treated as tissue
                                    luminosity_threshold=0.8,
                                    # herein we use 'ista' to compute the concentration
-                                   concentration_method='ista',
+                                   concentration_solver='ista',
                                    sigma_alpha=0.2,
                                    sigma_beta=0.2, target_stain_idx=(0, 1),
                                    # this allows to cache the stain matrix if it's too time-consuming to recompute.

demo.py

@@ -217,7 +217,7 @@ def postprocess(image_tensor): return convert_image_dtype(image_tensor, torch.ui
 fig, axs = plt.subplots(2, num_repeat + 1, figsize=(15, 8), dpi=300)
 for i, ax_alg in enumerate(axs):
     alg = algorithms[i].lower()
-    augmentor = AugmentorBuilder.build(alg, concentration_method='ista',
+    augmentor = AugmentorBuilder.build(alg, concentration_solver='ista',
                                        sigma_alpha=0.5,
                                        sigma_beta=0.5,
                                        luminosity_threshold=0.8,

tests/images/test_functionals.py

@@ -1,10 +1,15 @@
 """borrow Staintool's test cases
 """
 import unittest
+from typing import Optional, cast
+
 from tests.util import fix_seed, dummy_from_numpy, psnr
-from torch_staintools.functional.stain_extraction.macenko import MacenkoAlg
-from torch_staintools.functional.stain_extraction.vahadane import VahadaneAlg
-from torch_staintools.functional.concentration import get_concentrations
+from torch_staintools.constants import CONFIG
+from torch_staintools.functional.optimization.sparse_util import METHOD_FACTORIZE
+from torch_staintools.functional.stain_extraction.extractor import StainExtraction
+from torch_staintools.functional.stain_extraction.macenko import MacenkoAlg, DEFAULT_MACENKO_CONFIG
+from torch_staintools.functional.stain_extraction.vahadane import VahadaneAlg, DEFAULT_VAHADANE_CONFIG
+from torch_staintools.functional.concentration import ConcentrationSolver, ConcentCfg
 from torch_staintools.functional.tissue_mask import get_tissue_mask, TissueMaskException
 from torch_staintools.functional.utility.implementation import transpose_trailing, img_from_concentration
 from torchvision.transforms.functional import convert_image_dtype
@@ -24,6 +29,7 @@ class TestFunctional(unittest.TestCase):
     rand_img = torch.randint(0, 255, (1, 3, 256, 256))
 
     THRESH_PSNR = 20
+    POSITIVE_CONC_CFG = ConcentCfg()
 
     @staticmethod
     def get_dummy_path():
@@ -40,65 +46,87 @@ def new_dummy_img_tensor_ubyte():
         return TestFunctional.DUMMY_IMG_TENSOR.clone()
 
     @staticmethod
-    def stain_extract(dummy_tensor, get_stain_mat, luminosity_threshold, num_stains, algorithm, regularizer):
+    def stain_extract(dummy_tensor: torch.Tensor, get_stain_mat: StainExtraction,
+                      conc_solver: ConcentrationSolver,
+                      luminosity_threshold: float, num_stains: int, rng: Optional[torch.Generator]):
 
         # lab_tensor = rgb_to_lab(convert_image_dtype(dummy_tensor))
 
-        stain_matrix = get_stain_mat(image=dummy_tensor, luminosity_threshold=luminosity_threshold,
-                                     num_stains=num_stains, regularizer=regularizer)
+        stain_matrix = get_stain_mat(image=dummy_tensor,
+                                     luminosity_threshold=luminosity_threshold,
+                                     num_stains=num_stains, rng=rng)
 
-        concentration = get_concentrations(dummy_tensor, stain_matrix, algorithm=algorithm,
-                                           regularizer=regularizer)
+        concentration = conc_solver(dummy_tensor, stain_matrix, rng=rng)
         c_transposed_src = transpose_trailing(concentration)
         reconstructed = img_from_concentration(c_transposed_src, stain_matrix, dummy_tensor.shape, (0, 1))
         return stain_matrix, concentration, c_transposed_src, reconstructed
 
     @staticmethod
-    def extract_eval_helper(tester, get_stain_mat, luminosity_threshold,
-                            num_stains, regularizer, dict_algorithm):
+    def extract_eval_helper(tester, get_stain_mat: StainExtraction,
+                            conc_solver: ConcentrationSolver,
+                            luminosity_threshold: Optional[float],
+                            num_stains: int, rng: Optional[torch.Generator]):
         device = TestFunctional.device
         dummy_tensor_ubyte = TestFunctional.new_dummy_img_tensor_ubyte().to(device)
         # get_stain_mat = MacenkoExtractor()
         result_tuple = TestFunctional.stain_extract(dummy_tensor_ubyte, get_stain_mat,
+                                                    conc_solver=conc_solver,
                                                     luminosity_threshold=luminosity_threshold,
                                                     num_stains=num_stains,
-                                                    algorithm=dict_algorithm, regularizer=regularizer)
+                                                    rng=rng)
 
         stain_matrix, concentration, c_transposed_src, reconstructed = result_tuple
         dummy_scaled = convert_image_dtype(dummy_tensor_ubyte, torch.float32)
         psnr_out = psnr(dummy_scaled, reconstructed).item()
-        tester.assertTrue(psnr_out > TestFunctional.THRESH_PSNR)
+        tester.assertTrue(psnr_out > TestFunctional.THRESH_PSNR,
+                          msg=f"{psnr_out} vs. {TestFunctional.THRESH_PSNR}. \n"
+                              f"{get_stain_mat.stain_algorithm.cfg} \n"
+                              f"nan: {torch.isnan(reconstructed).any()} \n"
+                              f"Dict pos: {CONFIG.DICT_POSITIVE_DICTIONARY}")
         # size
         batch_size, channel_size, height, width = dummy_tensor_ubyte.shape
         tester.assertTrue(stain_matrix.shape == (batch_size, num_stains, channel_size))
 
         # transpose
         tester.assertTrue((c_transposed_src.permute(0, 2, 1) == concentration).all())
 
-        # manual tissue mask
-        mask = get_tissue_mask(dummy_scaled, luminosity_threshold=luminosity_threshold)
-        tissue_count = mask.sum()
-        tester.assertTrue(concentration.shape[-1] == tissue_count)
-
     def eval_wrapper(self, extractor):
 
         # all pixel
-        algorithms = ['ista', 'cd', 'ls']
-        for alg in algorithms:
-            TestFunctional.extract_eval_helper(self, extractor, luminosity_threshold=None,
-                                               num_stains=2, regularizer=0.1, dict_algorithm=alg)
+        algorithms = ['ista', 'cd', 'ls', 'fista']
+        dict_constraint_flag = [True]
+        for flag in dict_constraint_flag:
+            CONFIG.DICT_POSITIVE_DICTIONARY = flag
+            for alg in algorithms:
+                cfg = TestFunctional.POSITIVE_CONC_CFG
+                cfg.algorithm = cast(METHOD_FACTORIZE, alg)
+                cfg.positive = True
+                solver = ConcentrationSolver(cfg)
+                TestFunctional.extract_eval_helper(self, extractor, luminosity_threshold=None,
+                                                   num_stains=2, conc_solver=solver, rng=None)
+                solver.cfg.positive = False
+                TestFunctional.extract_eval_helper(self, extractor, luminosity_threshold=None,
+                                                   num_stains=2, conc_solver=solver, rng=None)
 
     def test_stains(self):
-        macenko = MacenkoAlg()
-        vahadane = VahadaneAlg()
+        macenko = StainExtraction(MacenkoAlg(DEFAULT_MACENKO_CONFIG))
+        vahadane = StainExtraction(VahadaneAlg(DEFAULT_VAHADANE_CONFIG))
         # not support num_stains other than 2
-        with self.assertRaises(NotImplementedError):
-            TestFunctional.extract_eval_helper(self, macenko, luminosity_threshold=None,
-                                               num_stains=3, regularizer=0.1, dict_algorithm='ista')
+        with self.assertRaises(AssertionError):
+            TestFunctional.extract_eval_helper(self, macenko,
+                                               conc_solver=ConcentrationSolver(TestFunctional.POSITIVE_CONC_CFG),
+                                               luminosity_threshold=None,
+                                               num_stains=3, rng=None)
 
         self.eval_wrapper(macenko)
         self.eval_wrapper(vahadane)
 
+        # vahadane with rng and lr
+        vahadane.stain_algorithm.cfg.lr = 0.5
+        TestFunctional.extract_eval_helper(self, vahadane,
+                                           conc_solver=ConcentrationSolver(TestFunctional.POSITIVE_CONC_CFG),
+                                           luminosity_threshold=None,
+                                           num_stains=3, rng=torch.Generator(1))
     def test_tissue_mask(self):
         device = TestFunctional.device
         dummy_scaled = convert_image_dtype(TestFunctional.new_dummy_img_tensor_ubyte(), torch.float32).to(device)

torch_staintools/augmentor/base.py

@@ -1,8 +1,8 @@
+from functools import partial
+
 import torch
 from typing import Optional, Sequence, Tuple, Hashable, List
-from ..functional.stain_extraction.factory import build_from_name
-from ..functional.optimization.sparse_util import METHOD_FACTORIZE
-from torch_staintools.functional.concentration import get_concentrations
+from torch_staintools.functional.concentration import ConcentrationSolver
 from ..functional.stain_extraction.extractor import StainExtraction, StainAlg
 from ..functional.utility.implementation import transpose_trailing, img_from_concentration
 from ..functional.tissue_mask import get_tissue_mask, TissueMaskException
@@ -23,10 +23,10 @@ class Augmentor(CachedRNGModule):
 
     # _tensor_cache: TensorCache
     # CACHE_FIELD: str = '_tensor_cache'
-
+    rng: Optional[torch.Generator]
     target_stain_idx: Optional[Sequence[int]]
 
-    concentration_method: METHOD_FACTORIZE
+    concentration_solver: ConcentrationSolver
     get_stain_matrix: StainExtraction  # can be any callable following the signature of BaseExtractor's __call__
     target_concentrations: torch.Tensor
 
@@ -38,7 +38,7 @@ class Augmentor(CachedRNGModule):
     regularizer: float
 
     def __init__(self, stain_alg: StainAlg,
-                 concentration_method: METHOD_FACTORIZE = 'ista',
+                 concentration_solver: ConcentrationSolver,
                  rng: TYPE_RNG = None,
                  target_stain_idx: Optional[Sequence[int]] = (0, 1),
                  sigma_alpha: float = 0.2,
@@ -56,7 +56,7 @@ def __init__(self, stain_alg: StainAlg,
         Args:
             stain_alg: the Callable to obtain stain matrix - e.g., Vahadane's dict learning or
                 Macenko's SVD
-            concentration_method:  How to get stain concentration from stain matrix
+            concentration_solver:  How to get stain concentration from stain matrix
             rng: the specified torch.Generator or int (as seed) for reproducing the results
             sigma_alpha: bound of alpha (mean 1). Sampled from (1-sigma, 1+sigma)
             sigma_beta: bound of beta (mean 0). Sampled from (-sigma, sigma)
@@ -67,7 +67,7 @@ def __init__(self, stain_alg: StainAlg,
 
         """
         super().__init__(cache, device, rng)
-        self.concentration_method = concentration_method
+        self.concentration_solver = concentration_solver
         self.get_stain_matrix = StainExtraction(stain_alg)
 
         self.target_stain_idx = target_stain_idx
@@ -212,32 +212,27 @@ def augment(*,
                                                                   alpha=alpha, beta=beta)
         return target_concentration
 
-    def forward(self, target: torch.Tensor, cache_keys: Optional[List[Hashable]] = None, **stain_mat_kwargs):
+    def forward(self, target: torch.Tensor, cache_keys: Optional[List[Hashable]] = None):
         """
 
         Args:
             target: input tensor to augment. Shape B x C x H x W and intensity range is [0, 1].
             cache_keys: unique keys point the input batch to the cached stain matrices. `None` means no cache.
-            **stain_mat_kwargs: all extra keyword arguments other than regularizer/num_stains/luminosity_threshold set
-                in __init__.
 
         Returns:
             Augmented output.
         """
         # stain_matrix_target -- B x num_stain x num_input_color_channel
         # todo cache
-        get_stain_mat_partial = self.get_stain_matrix.get_partial(luminosity_threshold=self.luminosity_threshold,
-                                                                  num_stains=self.num_stains,
-                                                                  regularizer=self.regularizer,
-                                                                  rng=self.rng,
-                                                                  **stain_mat_kwargs)
+        get_stain_partial = partial(self.get_stain_matrix,
+                                    luminosity_threshold=self.luminosity_threshold,
+                                    num_stains=self.num_stains, rng=self.rng)
 
-        target_stain_matrix = self.tensor_from_cache(cache_keys=cache_keys, func=get_stain_mat_partial,
+        target_stain_matrix = self.tensor_from_cache(cache_keys=cache_keys, func=get_stain_partial,
                                                      target=target)
 
         #  B x num_stains x num_pixel_in_mask
-        concentration = get_concentrations(target, target_stain_matrix, regularizer=self.regularizer,
-                                           algorithm=self.concentration_method, rng=self.rng)
+        concentration = self.concentration_solver(target, target_stain_matrix, rng=self.rng)
         try:
             tissue_mask = get_tissue_mask(target, luminosity_threshold=self.luminosity_threshold, throw_error=True,
                                           true_when_empty=False)
@@ -257,13 +252,14 @@ def forward(self, target: torch.Tensor, cache_keys: Optional[List[Hashable]] = N
 
     @classmethod
     def build(cls,
-              method: str, *, concentration_method: METHOD_FACTORIZE = 'ista',
+              stain_alg: StainAlg, *,
+              concentration_solver: ConcentrationSolver,
               rng: TYPE_RNG = None,
               target_stain_idx: Optional[Sequence[int]] = (0, 1),
               sigma_alpha: float = 0.2,
               sigma_beta: float = 0.2,
+              num_stains: int = 2,
               luminosity_threshold: Optional[float] = 0.8,
-              regularizer: float = 0.1,
               use_cache: bool = False,
               cache_size_limit: int = -1,
               device: Optional[torch.device] = None,
@@ -272,18 +268,18 @@ def build(cls,
         """Factory builder of the augmentor which manipulate the stain concentration by alpha * concentration + beta.
 
         Args:
-            method: algorithm name to extract stain - support 'vahadane' or 'macenko'
-            concentration_method: method to obtain the concentration. Default 'ista' for fast sparse solution on GPU
+            stain_alg: algorithm name to extract stain - support 'vahadane' or 'macenko'
+            concentration_solver: method to obtain the concentration. Default 'ista' for fast sparse solution on GPU
                 only applied for StainSeparation-based approaches (macenko and vahadane).
                 support 'ista', 'cd', and 'ls'. 'ls' simply solves the least square problem for factorization of
                 min||HExC - OD|| but is faster. 'ista'/cd enforce the sparse penalty but slower.
             rng: an optional seed (either an int or a torch.Generator) to determine the random number generation.
             target_stain_idx: what stains to augment: e.g., for HE cases, it can be either or both from [0, 1]
             sigma_alpha: alpha is uniformly randomly selected from (1-sigma_alpha, 1+sigma_alpha)
             sigma_beta: beta is uniformly randomly selected from (-sigma_beta, sigma_beta)
+            num_stains: number of stains to separate.
             luminosity_threshold: luminosity threshold to find tissue regions (smaller than but positive)
                 a pixel is considered as being tissue if the intensity falls in the open interval of (0, threshold).
-            regularizer: regularization term in ISTA algorithm
             use_cache: whether to use cache to save the stain matrix to avoid re-computation
             cache_size_limit: size limit of the cache. negative means no limits.
             device: what device to hold the cache.
@@ -293,11 +289,9 @@ def build(cls,
         Returns:
             Augmentor.
         """
-        method = method.lower()
-        extractor = build_from_name(method)
         cache = cls._init_cache(use_cache, cache_size_limit=cache_size_limit, device=device,
                                 load_path=load_path)
-        return cls(extractor, concentration_method=concentration_method, rng=rng, target_stain_idx=target_stain_idx,
-                   sigma_alpha=sigma_alpha, sigma_beta=sigma_beta,
-                   luminosity_threshold=luminosity_threshold, regularizer=regularizer,
+        return cls(stain_alg, concentration_solver=concentration_solver, rng=rng, target_stain_idx=target_stain_idx,
+                   sigma_alpha=sigma_alpha, sigma_beta=sigma_beta, num_stains=num_stains,
+                   luminosity_threshold=luminosity_threshold,
                    cache=cache, device=device).to(device)