torch_staintools.augmentor package

Submodules

torch_staintools.augmentor.base module

class torch_staintools.augmentor.base.Augmentor(stain_alg: StainAlg, concentration_solver: ConcentrationSolver, rng: int | Generator | None = None, target_stain_idx: Sequence[int] | None = (0, 1), sigma_alpha: float = 0.2, sigma_beta: float = 0.2, num_stains: int = 2, luminosity_threshold: float | None = 0.8, cache: TensorCache | None = None, device: device | None = None)

Bases: CachedRNGModule

Basic augmentation object as a nn.Module with stain matrices cache.

static augment(*, target_concentration: Tensor, tissue_mask: Tensor, target_stain_idx: Sequence[int] | None, inplace: bool, rng: Generator, sigma_alpha: float, sigma_beta: float)
Parameters:

  • target_concentration – concentration matrix of input image. B x num_stains x num_pixel

  • tissue_mask – region of the tissue

  • target_stain_idx – which stain channel to operate on.

  • inplace – whether augment the concentration matrix in-place

  • rng – rng for alpha and beta generation

  • sigma_alpha – sample values in range (-sigma, sigma)

  • sigma_beta – same semantic of sigma_alpha but applied to beta

Returns:

classmethod build(stain_alg: StainAlg, *, concentration_solver: ConcentrationSolver, rng: int | Generator | None = None, target_stain_idx: Sequence[int] | None = (0, 1), sigma_alpha: float = 0.2, sigma_beta: float = 0.2, num_stains: int = 2, luminosity_threshold: float | None = 0.8, use_cache: bool = False, cache_size_limit: int = -1, device: device | None = None, load_path: str | None = None)

Factory builder of the augmentor which manipulate the stain concentration by alpha * concentration + beta.

Parameters:

  • 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).

  • 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.

  • load_path – If specified, then stain matrix cache will be loaded from the file path. See the cache module for more details.

Returns:

Augmentor.

static channel_rand(target_concentration_selected: Tensor, rng: Generator, sigma_alpha: float, sigma_beta: float) → Tuple[Tensor, Tensor]
Parameters:

  • target_concentration_selected – concentrations to work on (e.g., the entire or a subset of concentration matrix)

  • rng – torch.Generator object

  • sigma_alpha – sample alpha values in range (1-sigma, 1+ sigma)

  • sigma_beta – sample beta values in range (-sigma, sigma)

Returns:

sampled alpha and beta as a tuple

concentration_solver: ConcentrationSolver
device: device
forward(target: Tensor, cache_keys: List[Hashable] | None = None)
Parameters:

  • 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.

Returns:

Augmented output.

get_stain_matrix: StainExtraction
luminosity_threshold: float
num_stains: int
static randn_range(*size, low, high, device: device, rng: Generator)

Helper function to get the uniform random float within low/high given the torch.Generator

Parameters:

  • *size – size of the random number

  • low – lower bound (inclusive)

  • high – upper bound (inclusive)

  • device – device to create the random numbers

  • rng – random number generator

Returns:

random sample given the size and the bounds using the specified rng.

regularizer: float
sigma_alpha: float
sigma_beta: float
target_concentrations: Tensor
target_stain_idx: Sequence[int] | None

torch_staintools.augmentor.factory module

class torch_staintools.augmentor.factory.AugmentorBuilder

Bases: object

Factory Builder for all supported normalizers: macenko/vahadane

For any non-stain separation-based augmentation, the factory builders can be integrated here for a unified interface.

static build(method: Literal['vahadane', 'macenko'], sparse_stain_solver: Literal['ista', 'cd', 'fista'] = 'fista', sparse_dict_steps: int = 60, dict_init: Literal['zero', 'transpose', 'unif', 'ridge'] = 'transpose', concentration_solver: Literal['ista', 'cd', 'fista', 'ls'] = 'fista', num_stains: int = 2, luminosity_threshold: float | None = 0.8, perc: int = 1, regularizer: float = 0.01, maxiter: int = 50, lr: float | None = None, tol: float = 1e-05, target_stain_idx: Sequence[int] | None = (0, 1), sigma_alpha: float = 0.2, sigma_beta: float = 0.2, use_cache: bool = False, cache_size_limit: int = -1, device: device | None = None, load_path: str | None = None, rng: int | Generator | None = None) → Augmentor

build from specified algorithm name method and augment the stain by alpha * concentration + beta

Warning

concentration_algorithm = ‘ls’ May fail on GPU for individual large input (e.g., 1000 x 1000), regardless of batch size. Therefore, ‘ls’ is better for multiple small inputs in terms of H and W.

Parameters:

  • method – Name of stain normalization algorithm. Support macenko and vahadane

  • sparse_stain_solver – sparse solver for dictionary learning in Vahadane algorithm. Support ‘ista’, ‘fista’, and ‘cd’.

  • sparse_dict_steps – steps of iteration in dictionary learning in Vahadane algorithm.

  • dict_init – code initialization method in dictionary learning in Vahadane algorithm.

  • concentration_solver – solver to obtain the concentration. Default ‘fista’ for fast sparse solution on GPU. Only applied to StainSeparation-based approaches (macenko and vahadane). support ‘fista’, ‘ista’, ‘cd’, and ‘ls’. ‘ls’ solves the least square problem for factorization of min||HExC - OD||, faster but no sparsity constraints. ‘ista’/cd enforce the sparse penalty but slower.

  • num_stains – number of stains to separate. Currently, Macenko only supports 2.

  • luminosity_threshold – luminosity threshold (smaller than) to find tissue region.

  • perc – Percentile threshold in Macenko algorithm to find the minimum angular term. min as 1 percentile and max angular as (100 - perc) percentile. Default is 1.

  • regularizer – regularization term in ISTA for dictionary learning (e.g., concentration computation)

  • maxiter – the max iteration in code updating of dictionary learning for stain matrix and concentration estimation (e.g., ISTA)

  • lr – learning rate for ISTA/FISTA-based optimization in code/concentration updating in ISTA/FISTA. If None, the invert of Lipschitz constant of the gradient is used.

  • tol – tolerance threshold for early convergence in ISTA/FISTA.

  • target_stain_idx – which stain to augment

  • sigma_alpha – alpha sampled from (1-sigma_alpha, 1+sigma_alpha)

  • sigma_beta – beta sampled from (-sigma_beta, sigma_beta)

  • use_cache – whether to cache the stain matrix for each input image

  • cache_size_limit – limit size of cache (how many matrices to cache). -1 means no limits.

  • device – device of the cache

  • load_path – whether to load prefetched cache. None means nothing will be loaded.

  • rng – random seed for augmentation and any random initialization may incur.

Returns:

corresponding Augmentor object.

Module contents