torch_staintools.augmentor package

Submodules

torch_staintools.augmentor.base module

class torch_staintools.augmentor.base.Augmentor(get_stain_matrix: BaseExtractor, concentration_method: Literal['ista', 'cd', 'ls'] = 'ista', 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, regularizer: float = 0.1, cache: TensorCache | None = None, device: device | None = None)

Bases: CachedRNGModule

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

__concentration_selected(target_stain_idx: Sequence[int] | None) → Tensor

Return concentration of selected stain channels

Parameters:

  • target_concentration – B x num_stains x num_pixel_in_mask

  • target_stain_idx – Basic indices of the selected stains. If None then returns all stain.

Returns:

The view of the selected stain (no copies)

__inplace_augment_helper(*, tissue_mask: Tensor, alpha: Tensor, beta: Tensor)

Helper function to augment a given row(s) of stain concentration: alpha * concentration + beta

Parameters:

  • target_concentration – B x num_stains x num_pixel

  • tissue_mask – mask of tissue regions. only augment concentration within the mask

  • alpha – alpha value for augmentation

  • beta – beta value for augmentation

Returns:

augmented concentration. The result is in-place (as for the target_concentration here). It might be cloned beforehand in prior.

__inplace_tensor(inplace: bool) → Tensor

Helper function to get the clone or the original concentration.

Parameters:

  • target_concentration – concentration tensor

  • inplace – bool. If True then returns itself, otherwise returns a clone.

Returns:

The original or cloned concentration tensor

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(method: str, *, concentration_method: Literal['ista', 'cd', 'ls'] = 'ista', rng: int | Generator | None = None, target_stain_idx: Sequence[int] | None = (0, 1), sigma_alpha: float = 0.2, sigma_beta: float = 0.2, luminosity_threshold: float | None = 0.8, regularizer: float = 0.1, 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:

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

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

  • 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_method: Literal['ista', 'cd', 'ls']
device: device
forward(target: Tensor, cache_keys: List[Hashable] | None = None, **stain_mat_kwargs)
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.

  • **stain_mat_kwargs – all extra keyword arguments other than regularizer/num_stains/luminosity_threshold set in __init__.

Returns:

Augmented output.

get_stain_matrix: BaseExtractor
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'], concentration_method: Literal['ista', 'cd', 'ls'] = 'ista', rng: int | Generator | None = None, target_stain_idx: Sequence[int] | None = (0, 1), sigma_alpha: float = 0.2, sigma_beta: float = 0.2, luminosity_threshold: float | None = 0.8, regularizer: float = 0.1, use_cache: bool = False, cache_size_limit: int = -1, device: device | None = None, load_path: str | 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

  • concentration_method – 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 – random seed for augmentation and any random initialization may incur.

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

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

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

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

Returns:

corresponding Augmentor object.

Module contents