Package: torch-staintools
Torch StainTools for Stain Normalization and Augmentation of Histopathological ImagesInstallationFrom Repository:pip install git+https://github.com/CielAl/torch-staintools.gitFrom PyPI:pip install torch-staintoolsDocumentationDetail documentation regarding the code base can be found in theGitPages.CitationIf this toolkit helps you in your publication, please feel free to cite with the following bibtex entry:@software{zhou_2024_10453807,author={Zhou, Yufei},title={CielAl/torch-staintools: V1.0.3 Release},month=jan,year=2024,publisher={Zenodo},version={v1.0.3},doi={10.5281/zenodo.10453807},url={https://doi.org/10.5281/zenodo.10453807}}DescriptionStain Normalization (Reinhard, Macenko, and Vahadane) for pytorch. Input tensors (fit and transform) must be in shape ofNxCxHxW, with value scaled to [0, 1] in format of torch.float32.Stain Augmentation using Macenko and Vahadane as stain extraction.Fast normalization/augmentation on GPU with stain matrices caching.Simulate the workflow inStainTools librarybut use the Iterative Shrinkage Thresholding Algorithm (ISTA), or optionally, the coordinate descent (CD) to solve the dictionary learning for stain matrix computation in Vahadane or Macenko (stain concentration only) algorithm. The implementation of ISTA and CD are derived from Cédric Walker'storchvahadaneStain Concentration is solved via factorization ofStain_Matrix x Concentration = Optical_Density. For efficient sparse solution and more robust outcomes, ISTA can be applied. Alternatively, Least Square solver (LS) fromtorch.linalg.lstsqmight be applied for faster non-sparse solution.No SPAMS requirement (which is a dependency in StainTools).Sample Output of Torch-StainTools NormalizationSample Output of StainToolsSample Output of Torch-StainTools Augmentation (Repeat 3 times)Sample Output of StainTools Augmentation (Repeat 3 times)Benchmark (No Stain Matrices Caching)Use the sample images under ./test_images (size2500x2500x3). Mean was computed from 7 runs (1 loop per run) using
timeit. Comparison between torch_stain_tools in CPU/GPU mode, as well as that of the StainTools Implementation.For consistency, use ISTA to compute the concentration.TransformationMethodCPU[s]GPU[s]StainTool[s]Vahadane1197.520.9Macenko5.570.47920.7Reinhard0.8400.0240.414FittingMethodCPU[s]GPU[s]StainTool[s]Vahadane1328.4019.1Macenko6.990.06420.0Reinhard0.4220.0110.076Batchified Concentration ComputationSplit the sample images under ./test_images (size2500x2500x3) into 81 non-overlapping256x256x3tiles as a batch.For the StainTools baseline, a for-loop is implemented to get the individual concentration of each of the numpy array of the 81 tiles.MethodCPU[s]GPU[s]ISTA (concentration_method='ista')3.121.24CD   (concentration_method='cd')29.3s4.87LS   (concentration_method='ls')0.2210.097StainTools (SPAMS)16.6N/AUse Cases and TipsFor details, follow the example in demo.pyNormalizers are wrapped astorch.nn.Module, working similarly to a standalone neural network. This means that for a workflow involving dataloader with multiprocessing, the normalizer
(Note that CUDA has poor support in multiprocessing, and therefore it may not be the best practice to perform GPU-accelerated on-the-fly stain transformation in pytorch's dataset/dataloader)concentration_method='ls'(i.e.,torch.linalg.lstsq) can be efficient for batches of many smaller input (e.g.,256x256) in terms of width and height. However, it may fail on GPU for a single larger input image (width and height). This happens even if the
the total number of pixels of the image is fewer than the aforementioned batch of multiple smaller input. Therefore,concentration_method='ls'could be suitable to deal with huge amount of small images in batches on the fly.importcv2importtorchfromtorchvision.transformsimportToTensorfromtorch_staintools.normalizerimportNormalizerBuilderfromtorch_staintools.augmentorimportAugmentorBuilderimportosseed=0torch.manual_seed(seed)torch.cuda.manual_seed(seed)# cpu or gpudevice=torch.device("cuda:0")root_dir='.'target=cv2.imread(os.path.join(root_dir,'test_images/TCGA-33-4547-01Z-00-DX7.''91be6f90-d9ab-4345-a3bd-91805d9761b9_8270_5932_0.png'))# shape: Height (H) x Width (W) x Channel (C, for RGB C=3)target=cv2.cvtColor(target,cv2.COLOR_BGR2RGB)norm=cv2.imread(os.path.join(root_dir,'test_images/TCGA-95-8494-01Z-00-DX1.''716299EF-71BB-4095-8F4D-F0C2252CE594_5932_5708_0.png'))# shape: HWCnorm=cv2.cvtColor(norm,cv2.COLOR_BGR2RGB)# shape: Batch x Channel x Height x Width (BCHW); in the showcase here batch size is 1 (B=1) - scaled to [0, 1] torch.float32target_tensor=ToTensor()(target).unsqueeze(0).to(device)# shape: BCHW - scaled to [0, 1] torch.float32norm_tensor=ToTensor()(norm).unsqueeze(0).to(device)# ######## Normalization# create the normalizer - using vahadane. Alternatively can use 'macenko' or 'reinhard'.# note this is equivalent to:# from torch_staintools.normalizer.separation import StainSeparation# normalizer_vahadane = StainSeparation.build('vahadane', **arguments)# we use the 'ista' (ISTA algorithm) to get the sparse solution of the factorization: STAIN_MATRIX * Concentration = OD# alternatively, 'cd' (coordinate descent) and 'ls' (least square from torch.linalg) is available.# Note that 'ls' does not can be much faster on batches of smaller input, but may fail on GPU for individual large input# in terms of width and height, regardless of the batch sizenormalizer_vahadane=NormalizerBuilder.build('vahadane',concentration_method='ista')# move the normalizer to the device (CPU or GPU)normalizer_vahadane=normalizer_vahadane.to(device)# fit. For macenko and vahadane this step will compute the stain matrix and concentrationnormalizer_vahadane.fit(target_tensor)# transform# BCHW - scaled to [0, 1] torch.float32output=normalizer_vahadane(norm_tensor)# ###### Augmentation# augment by: alpha * concentration + beta, while alpha is uniformly randomly sampled from (1 - sigma_alpha, 1 + sigma_alpha),# and beta is uniformly randomly sampled from (-sigma_beta, sigma_beta).augmentor=AugmentorBuilder.build('vahadane',# fix the random number generator seed for reproducibility.rng=314159,# the luminosity threshold to find the tissue region to augment# if set to None means all pixels are treated as tissueluminosity_threshold=0.8,# herein we use 'ista' to compute the concentrationconcentration_method='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.# e.g., if using Vahadane algorithmuse_cache=True,# size limit of cache. -1 means no limit (stain matrix is often small in size, e.g., 2 x 3)cache_size_limit=-1,# if specified, the augmentor will load the cached stain matrices from file system.load_path=None,)# move augmentor to the corresponding deviceaugmentor=augmentor.to(device)num_augment=5# multiple copies of different random augmentation of the same tile may be generatedfor_inrange(num_augment):# B x C x H x W# use a list of Hashable key (e.g., str) to map the batch input to its corresponding stain matrix in cache.# this key should be unique, e.g., using the filename of the input tile.# leave it as None if no caching is intended, even if use_cache is enabled.# note since the inputs are all batchified, the cache_key are in form of a list, with each element in the# list corresponding to a data point in the batch.aug_out=augmentor(norm_tensor,cache_keys=['some unique key'])# do anything to the augmentation output# dump the cache of stain matrices for future usageaugmentor.dump_cache('./cache.pickle')# fast batch operationtile_size=512tiles:torch.Tensor=norm_tensor.unfold(2,tile_size,tile_size)\.unfold(3,tile_size,tile_size).reshape(1,3,-1,tile_size,tile_size).squeeze(0).permute(1,0,2,3).contiguous()print(tiles.shape)# use macenko normalization as examplenormalizer_macenko=NormalizerBuilder.build('macenko',use_cache=True,# use least square solver, along with cache, to perform# normalization on-the-flyconcentration_method='ls')normalizer_macenko=normalizer_macenko.to(device)normalizer_macenko.fit(target_tensor)normalizer_macenko(tiles)Stain Matrix CachingAs elaborated in the below in the running time benchmark of fitting, computation of stain matrix could be time-consuming.
Therefore, for bothAugmentorandNormalizer, an in-memory (device-specified) cache is implemented to store the previously computed stain matrices (typically with size 2 x 3 in H&E/RGB cases).
To enable the feature, theuse_cachemust be enabled, should you use the factory builders to instantiate theNormalizerorAugmentor.
Upon the normalization/augmentation procedure, a unique cache_key corresponding to the image input must be defined (e.g., file name).
Since bothNormalizerandAugmentorare designed astorch.nn.Moduleto accept batch inputs (tensors of shape B x C x H x W), a list of cache_keys must be given along with the batch image
inputs during the forward passing:normalizer_vahadane(input_batch, cache_keys=list_of_keys_corresponding_to_input_batch)
augmentor(input_batch, cache_keys=list_of_keys_corresponding_to_input_batch)The next timeNormalizerorAugmentorprocess the images, the corresponding stain matrices will be queried and fetched from cache if they are stored already, rather than recomputing from scratch.AcknowledgmentsSome codes are derived fromtorchvahadane,torchstain, andStainToolsSample images in the demo and ReadMe.md are selected fromThe Cancer Genome Atlas Program(TCGA)dataset.