MONAI framework is very slow while computing stain normalisation

[VinayMushunuri]

Hello,

I am new to MONAI framework and I am working on stain normalization approach to generate images similar to target image using MONAI framework.
I have checked the blog for the implementation of Stain normalization and I found out that the implementation is using numpy rather than cupy or something similar.

Here is the implementation, I found from MONAI blog :
Copyright (c) MONAI Consortium
Licensed under the Apache License, Version 2.0 (the “License”);
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an “AS IS” BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

from future import annotations

import numpy as np

from monai.transforms.transform import Transform

[docs]
class ExtractHEStains(Transform):
“”"Class to extract a target stain from an image, using stain deconvolution (see Note).

Args:
tli: transmitted light intensity. Defaults to 240.
alpha: tolerance in percentile for the pseudo-min (alpha percentile)
and pseudo-max (100 - alpha percentile). Defaults to 1.
beta: absorbance threshold for transparent pixels. Defaults to 0.15
max_cref: reference maximum stain concentrations for Hematoxylin & Eosin (H&E).
Defaults to (1.9705, 1.0308).

Note:
For more information refer to:
- the original paper: Macenko et al., 2009 http://wwwx.cs.unc.edu/~mn/sites/default/files/macenko2009.pdf
- the previous implementations:

  - MATLAB: https://github.com/mitkovetta/staining-normalization
  - Python: https://github.com/schaugf/HEnorm_python

"""

def init(
self, tli: float = 240, alpha: float = 1, beta: float = 0.15, max_cref: tuple | np.ndarray = (1.9705, 1.0308)
) -> None:
self.tli = tli
self.alpha = alpha
self.beta = beta
self.max_cref = np.array(max_cref)

def _deconvolution_extract_stain(self, image: np.ndarray) -> np.ndarray:
"""Perform Stain Deconvolution and return stain matrix for the image.

Args:
    image: uint8 RGB image to perform stain deconvolution on

Return:
    he: H&E absorbance matrix for the image (first column is H, second column is E, rows are RGB values)
"""
# check image type and values
if not isinstance(image, np.ndarray):
    raise TypeError("Image must be of type numpy.ndarray.")
if image.min() < 0:
    raise ValueError("Image should not have negative values.")
if image.max() > 255:
    raise ValueError("Image should not have values greater than 255.")

# reshape image and calculate absorbance
image = image.reshape((-1, 3))
image = image.astype(np.float32, copy=False) + 1.0
absorbance = -np.log(image.clip(max=self.tli) / self.tli)

# remove transparent pixels
absorbance_hat = absorbance[np.all(absorbance > self.beta, axis=1)]
if len(absorbance_hat) == 0:
    raise ValueError("All pixels of the input image are below the absorbance threshold.")

# compute eigenvectors
_, eigvecs = np.linalg.eigh(np.cov(absorbance_hat.T).astype(np.float32, copy=False))

# project on the plane spanned by the eigenvectors corresponding to the two largest eigenvalues
t_hat = absorbance_hat.dot(eigvecs[:, 1:3])

# find the min and max vectors and project back to absorbance space
phi = np.arctan2(t_hat[:, 1], t_hat[:, 0])
min_phi = np.percentile(phi, self.alpha)
max_phi = np.percentile(phi, 100 - self.alpha)
v_min = eigvecs[:, 1:3].dot(np.array([(np.cos(min_phi), np.sin(min_phi))], dtype=np.float32).T)
v_max = eigvecs[:, 1:3].dot(np.array([(np.cos(max_phi), np.sin(max_phi))], dtype=np.float32).T)

# a heuristic to make the vector corresponding to hematoxylin first and the one corresponding to eosin second
if v_min[0] > v_max[0]:
    he = np.array((v_min[:, 0], v_max[:, 0]), dtype=np.float32).T
else:
    he = np.array((v_max[:, 0], v_min[:, 0]), dtype=np.float32).T

return he

def call(self, image: np.ndarray) -> np.ndarray:
"""Perform stain extraction.

Args:
    image: uint8 RGB image to extract stain from

return:
    target_he: H&E absorbance matrix for the image (first column is H, second column is E, rows are RGB values)
"""
if not isinstance(image, np.ndarray):
    raise TypeError("Image must be of type numpy.ndarray.")

target_he = self._deconvolution_extract_stain(image)
return target_he

[docs]
class NormalizeHEStains(Transform):
“”"Class to normalize patches/images to a reference or target image stain (see Note).

Performs stain deconvolution of the source image using the ExtractHEStains
class, to obtain the stain matrix and calculate the stain concentration matrix
for the image. Then, performs the inverse Beer-Lambert transform to recreate the
patch using the target H&E stain matrix provided. If no target stain provided, a default
reference stain is used. Similarly, if no maximum stain concentrations are provided, a
reference maximum stain concentrations matrix is used.

Args:
tli: transmitted light intensity. Defaults to 240.
alpha: tolerance in percentile for the pseudo-min (alpha percentile) and
pseudo-max (100 - alpha percentile). Defaults to 1.
beta: absorbance threshold for transparent pixels. Defaults to 0.15.
target_he: target stain matrix. Defaults to ((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)).
max_cref: reference maximum stain concentrations for Hematoxylin & Eosin (H&E).
Defaults to [1.9705, 1.0308].

Note:
For more information refer to:
- the original paper: Macenko et al., 2009 http://wwwx.cs.unc.edu/~mn/sites/default/files/macenko2009.pdf
- the previous implementations:

    - MATLAB: https://github.com/mitkovetta/staining-normalization
    - Python: https://github.com/schaugf/HEnorm_python

"""

def init(
self,
tli: float = 240,
alpha: float = 1,
beta: float = 0.15,
target_he: tuple | np.ndarray = ((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)),
max_cref: tuple | np.ndarray = (1.9705, 1.0308),
) -> None:
self.tli = tli
self.target_he = np.array(target_he)
self.max_cref = np.array(max_cref)
self.stain_extractor = ExtractHEStains(tli=self.tli, alpha=alpha, beta=beta, max_cref=self.max_cref)

def call(self, image: np.ndarray) -> np.ndarray:
"""Perform stain normalization.

Args:
    image: uint8 RGB image/patch to be stain normalized, pixel values between 0 and 255

Return:
    image_norm: stain normalized image/patch
"""
# check image type and values
if not isinstance(image, np.ndarray):
    raise TypeError("Image must be of type numpy.ndarray.")
if image.min() < 0:
    raise ValueError("Image should not have negative values.")
if image.max() > 255:
    raise ValueError("Image should not have values greater than 255.")

# extract stain of the image
he = self.stain_extractor(image)

# reshape image and calculate absorbance
h, w, _ = image.shape
image = image.reshape((-1, 3))
image = image.astype(np.float32) + 1.0
absorbance = -np.log(image.clip(max=self.tli) / self.tli)

# rows correspond to channels (RGB), columns to absorbance values
y = np.reshape(absorbance, (-1, 3)).T

# determine concentrations of the individual stains
conc = np.linalg.lstsq(he, y, rcond=None)[0]

# normalize stain concentrations
max_conc = np.asarray([np.percentile(conc[0, :], 99), np.percentile(conc[1, :], 99)], dtype=np.float32)
tmp = np.divide(max_conc, self.max_cref, dtype=np.float32)
image_c = np.divide(conc, tmp[:, np.newaxis], dtype=np.float32)

image_norm: np.ndarray = np.multiply(self.tli, np.exp(-self.target_he.dot(image_c)), dtype=np.float32)
image_norm[image_norm > 255] = 254
image_norm = np.reshape(image_norm.T, (h, w, 3)).astype(np.uint8)
return image_norm

For me when I am running on a folder which contains 5400 webp tiles it is taking around half an hour. I have more than 2000 such folders and it takes at least a month to complete this procedure.

Could u please let me know if there is anything I could do to speed up the process.

[KumoLiu]

Hi @drbeh, could you please help take a look at this issue?
Thanks in advance!

[drbeh]

Hi @VinayMushunuri,

You are right, this implementation is based on numpy but we do have a cupy implementation that is way faster than numpy.
For you to use the cupy-based version, there are two options:

1. We have implemented a cupy based in Refactor Stain Normalization #3505 a while ago but we decided no to have the cupy version in MONAI! So you can checkout the PR and use it if you want.
2. The other option, which is the preferred one, is to use what we have implemented in cuCIM.
   You can use them directly or use it in MONAI via CuCIM or CuCIMd wrapper.

from tranforms import CuCIM
normalizer = CuCIM(
    name="normalize_colors_pca",
    source_intensity=240.0,
    alpha=1.0,
    beta=0.345,
    ref_stain_coeff=((0.5626, 0.2159), (0.7201, 0.8012), (0.4062, 0.5581)),
    ref_max_conc=(1.9705, 1.0308),
    image_type="intensity",
    channel_axis=0,
)
normalized_image = normalizer(image)

or simply:

from tranforms import CuCIM
normalizer = CuCIM(name="normalize_colors_pca")
normalized_image = normalizer(image)

[zhanghaoyue]

The call function from CuCIM is changed, needs an easy change from Monai side to fix the compatibility issue

[VinayMushunuri]

Hello @drbeh,
Thankyou for the response I would check it very quickly and let u know how well it works

[gigony]

Hi @VinayMushunuri,

Sending tiles to GPU takes some time indeed. May I ask what is the GPU that you are using?

I would suggest to profile your code to know how much of time is spent in loading and how much in stain normalization.

@gigony @grlee77 are you aware of a way to directly load webp tiles into gpu by any chance?

Hi @drbeh , cuCIM doesn't support webp compression or webp decoder for now, and I don't know any CUDA implementations available for WebP.

[VinayMushunuri]

I just checked the time it takes. Surprisingly, it is taking same time as MONAI to perform normalization and it is taking 0.2 sec just to put it in cuda.

[VinayMushunuri]

will next release include support for webp files or is it not gonna be part of it?

[drbeh]

We don't have the support of WebP on the roadmap for now but can you provide more information on your data/setup so that we can help to understand why you don't get the acceleration? Like

• what is the size of each tile?
• how are your files stored? Are they stored as individual tiles and you load each of the separately using PIL?
• how many tiles are you dealing with?
• what is the gpu that you are running on? or can you provide the output of the following code:

python -c 'import monai; monai.config.print_debug_info()'

[VinayMushunuri]

1. size of each tile is 962x962x3
2. Files are stored as individual tiles and each tiles are loaded individually using PIL
3. Each case or subject contains minimum of 5000 and maximum of 12000 tiles
4. I am running gpu titan

Output of following command
Printing MONAI config...

MONAI version: 0.8.1
Numpy version: 1.19.5
Pytorch version: 1.6.0
MONAI flags: HAS_EXT = False, USE_COMPILED = False
MONAI rev id: 71ff399

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: NOT INSTALLED or UNKNOWN VERSION.
scikit-image version: 0.17.2
Pillow version: 7.2.0
Tensorboard version: 2.5.0
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.7.0
tqdm version: 4.64.1
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: 5.8.0
pandas version: 1.1.1
einops version: NOT INSTALLED or UNKNOWN VERSION.
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.

For details about installing the optional dependencies, please visit:
https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies

================================
Printing system config...

System: Linux
Linux version: Ubuntu 18.04.5 LTS
Platform: Linux-5.4.0-137-generic-x86_64-with-Ubuntu-18.04-bionic
Processor: x86_64
Machine: x86_64
Python version: 3.6.9
Process name: python
Command: ['python', '-c', 'import monai; monai.config.print_debug_info()']
Open files: []
Num physical CPUs: 4
Num logical CPUs: 8
Num usable CPUs: 8
CPU usage (%): [13.2, 72.0, 8.4, 10.1, 7.5, 7.5, 9.2, 35.5]
CPU freq. (MHz): 1224
Load avg. in last 1, 5, 15 mins (%): [10.8, 8.9, 5.5]
Disk usage (%): 98.5
Avg. sensor temp. (Celsius): UNKNOWN for given OS
Total physical memory (GB): 31.3
Available memory (GB): 25.5
Used memory (GB): 5.2

================================
Printing GPU config...

Num GPUs: 1
Has CUDA: True
CUDA version: 10.2
cuDNN enabled: True
cuDNN version: 7605
Current device: 0
Library compiled for CUDA architectures: ['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75']
GPU 0 Name: GeForce GTX TITAN X
GPU 0 Is integrated: False
GPU 0 Is multi GPU board: False
GPU 0 Multi processor count: 24
GPU 0 Total memory (GB): 11.9
GPU 0 CUDA capability (maj.min): 5.2

[VinayMushunuri]

Also I found a bug in cucim.core.operations.color.stain_normalizer

In line 123 and 131 it is returning cp.round(image) but ccording to cupy the correct syntax is cp.round_(image). It threw an error here for me, so please find this once.

[gigony]

Thanks @VinayMushunuri !
@grlee77 , could you please check the code?

[drbeh]

That's right, cupy.around should be used instead.

[grlee77]

Thanks for pointing this out, @VinayMushunuri. This issue should not occur for CuPy>=11, but we should indeed switch from round to around so that we are also compatible with CuPy 10.x releases.

I opened a PR to fix that here: rapidsai/cucim#508. We are already in code freeze for the Feb 2023 version of cuCIM, but that fix should be available in the April 2023 release.

[grlee77]

A couple of questions regarding performance of the cuCIM implementation:
1.) what is the shape and dtype of the image being used here?
2.) can you provide the output of the following command:
python -c "import cupy; cupy.show_config()"

Assuming you are using CuPy<11 it may help to also define the environment variable: CUPY_ACCELERATORS="cub"
CUB is an optional library that is used to accelerate some types of numerical operations in CuPy. In CuPy>=11, CUB should already be enabled by default, so setting the environment variable should not be necessary.

[VinayMushunuri]

1. Shape of the image is 962x962x3 and dtype is

2. The out of following command is this:

OS : Linux-5.4.0-137-generic-x86_64-with-Ubuntu-18.04-bionic
CuPy Version : 8.5.0
NumPy Version : 1.19.5
SciPy Version : 1.4.1
Cython Build Version : 0.29.22
CUDA Root : /usr/local/cuda
CUDA Build Version : 10010
CUDA Driver Version : 11000
CUDA Runtime Version : 10010
cuBLAS Version : 10201
cuFFT Version : 10101
cuRAND Version : 10101
cuSOLVER Version : (10, 2, 0)
cuSPARSE Version : 10300
NVRTC Version : (10, 1)
Thrust Version : 100906
CUB Build Version : 100800
cuDNN Build Version : 8005
cuDNN Version : 8003
NCCL Build Version : 2803
NCCL Runtime Version : 2803
cuTENSOR Version : None
Device 0 Name : GeForce GTX TITAN X
Device 0 Compute Capability : 52

Since my cupy version is less than 11, I added this command os.environ["CUPY_ACCELERATORS"]="cub". But it still didn't improve the performance. It is still the same..