Ultra-high-resolution unpaired stain transformation via Kernelized Instance Normalization

Several frameworks have been proposed for unpaired image-to-image translation. CycleGAN [33], DiscoGAN [19], and DualGAN [31] were first presented to overcome the supervised pairing constraint via cycle consistency. However, subtle information was forced to be retained in the translated image to achieve better reconstruction, causing detrimental effects when two domains are substantially different such as dog-to-cat translation. Besides, a reverse mapping function might not always exist, which inevitably leads to artifacts in translated images. Recently, strategies beyond cyclic loss have been developed to reach one-sided unpaired image-to-image translation. While DistanceGAN [3] enforced distance consistency between different parts of the same sample in each domain, CUT [28] leveraged contrastive loss to maximize the patch-wise similarity between domains and achieved remarkable results. LSeSim [32] further utilized spatially correlation to maximize structural similarity and eliminate the domain-specific features.

Transforming one stained tissue into another specific stain will dramatically save laboratory resources and money. Hence, growing research has leveraged unsupervised image-to-image translation to conduct stain transformation in several medical scenarios. Levy et al. . translated H&E to trichrome staining via CycleGAN for liver fibrosis staging [22]. Kapil et al. translated Cytokeratin to PD-L1 staining to bypass re-staining for segmentation [18]. de Haan et al. translated H&E to Masson’s Trichrome, periodic acid-Schiff, and Jones silver stain for improving preliminary diagnosis for kidney diseases via CycleGAN with perfectly paired images [13]. Lahiani et al. further broke the limitation of $256\times 256$-pixel image patches by applying perceptual embedding consistency for H&E to FAP-CK transformation [21]. However, the lack of detailed description hampers the implementation of their methodology.

Ultra-high-resolution images are ubiquitous in photography, artwork, posters, ultra-high (e.g., 8K) videos, and especially, WSIs in digital pathology (usually, larger than $10,000\times 10,000$). Due to the massive computational costs, transforming these images into different styles will be difficult. Traditionally, strategies have been proposed to address the problem of tiling artifacts created by patch-wise-based methods, including utilizing a larger overlapping window [20] or freezing IN layers during testing time [2]. By providing a patch-wise style transfer network with Thumbnail Instance Normalization (TIN), Chen et al. performed ultra-high-resolution style transformation with constant GPU memory usage [8]. Also, they applied their framework to an image-to-image translation task. However, according to our experiment, TIN may result in over/under-colorizing for their fallacious assumption that all patches can be normalized with the same global mean and standard deviation.

Our framework targets one-sidedly translating an ultra-high-resolution image $X$ in domain $\mathcal{X}$ (e.g., H&E stain domain) into image $\hat{Y}$ in domain $\mathcal{Y}$ (e.g., IHC domain), in which $X,\hat{Y}\in R^{H\times W\times C}$, $H$ and $W$ are the height and width of $X$, via a mapping function, generator $\mathcal{G}$.

Collections of unpaired X in $\mathcal{X}$ and Y in $\mathcal{Y}$ would be first cropped into patches with the size of $512\times 512$ pixels to train a generator $\mathcal{G}$.

As our KIN module is only applied during the testing time and can be inserted into any framework with IN layers, we followed the original training process and hyperparameters proposed in the paper of CycleGAN [33], CUT [28], and LSeSim [32] to train the corresponding generators with their specific designed losses without any modification.

During the testing process, all the IN layers in $\mathcal{G}$ are replaced with KIN layers. Given an image $X$, non-overlapped patches $x^{i,j}_{p}$ are cropped with the size of $512\times 512$. The coordinates $i,j$ of each patch $x^{i,j}_{p}$ corresponding to the original $X$ would be recorded simultaneously. For example, an $M\times N$ image would be cropped into $\lfloor M/512\rfloor\times\lfloor N/512\rfloor$ patches with coordinates of {$0,1,...,\lfloor M/512\rfloor$} $\times$ {$0,1,...,\lfloor N/512\rfloor$}. Two caching tables of size $\lfloor M/512\rfloor\times\lfloor N/512\rfloor\times C$, in which $C$ denotes the number of channels, would be initialized in each KIN for caching mean and standard deviation calculated.

As illustrated in Fig 2, we divide the testing process into two phases: caching and inference. During caching phase, each patch ${x^{i,j}_{p}}$ with its corresponding coordinates $i$, $j$ are the input of the generator $\mathcal{G}$, and the calculated mean $\mu(x^{i,j}_{p})$ and standard deviation $\sigma(x^{i,j}_{p})$ after passing the KIN will be cached. During the inference phase, ${x^{i,j}_{p}}$, its corresponding coordinates $i$, $j$ and a kernel $k$ are the input of the generator $\mathcal{G}$. The kernel $k\in R^{h\times w}$, where $h$ and $w$ are the height and width of $k$, is adjustable. When passing through the KIN layer, a region with the same size of kernel $k$ extended from ${i,j}$ will be extracted from the caching table and convolute with kernel k to compute mean $\mu_{KIN}(x^{i,j}_{p})$ and standard deviation $\sigma_{KIN}(x^{i,j}_{p})$ which are used to normalize the feature maps. All the cropped patches will be passed to the $\mathcal{G}$ in the aforementioned manner to yield translated patches $\hat{y}^{i,j}_{p}$. Eventually, all $\hat{y}^{i,j}_{p}$ are assembled into an ultra-high-resolution translated image $\hat{Y}$.

IN [29] has been widely used in GAN-based models for image generation and dramatically improved image quality [27]. Besides, multiple styles can be obtained by conditionally replacing the $\mu$ and $\sigma$ in the IN layer [9]. IN can be formulated by:

For each instance in a batch, $\mu(X)$ and $\sigma(X)$ are calculated in a channel-wise manner, in which $\mu(X),\sigma(X)\in R^{B\times C}$, $B$ denotes the batch size and $\gamma$ and $\beta$ are trainable parameters.

We hypothesize that adjacent patches share similar statistics including the $\mu$ and $\sigma$ computed in IN, and thus proposed KIN that could further alleviate the subtle incongruity that induces the tiling artifacts when adjacent patches are assembled. KIN is the extension of the original IN layer with extra two caching tables $T_{\mu}$ and $T_{\sigma}$ to spatially store $\mu(X)$ and $\sigma(X)$ values and additionally supports convolution operation on the caching tables with a given kernel $k$. During the caching phase, KIN input a cropped patch $x^{i,j}_{p}$ with its spatial information, $i,j$. $\mu(x^{i,j}_{p})$ and $\sigma(x^{i,j}_{p})$ are computed as the original IN and cached.

During the inference phase, given a kernel $k$ with the size of $2q+1$, $\mu_{KIN}(x^{i,j}_{p})$ and $\sigma_{KIN}(x^{i,j}_{p})$ are computed by convoluting $k$ on cache tables to generate translated images. To address the boundary cases, the cache tables would be padded initially with edge values.

Automatic Non-rigid Histological Image Registration (ANHIR) dataset [5, 6, 7, 10, 12, 24] consists of high-resolution WSIs from different tissue samples (lesions, lung lobes, breast tissue, kidney tissue, gastric tissue, colon adenocarcinoma, and mammary gland). The acquired images are organized in sets of consecutive tissue slices stained by various dyes, including H&E, Ki-67, ER/PR, CD4/CD8/CD68, etc., with sizes vary from $15,000\times 15,000$ to $50,000\times 50,000$ pixels. We randomly sampled three types of tissues to conduct our experiments. Each experiment comprises H&E stain and one target IHC stain: breast tissue (from H&E to PR), colon adenocarcinoma (COAD) (from H&E to CD4&CD68), and lung lesion (from H&E to Ki-67).

The private glioma dataset was collected from H&E ($98,304\times 93,184$ pixels) and epidermal growth factor receptor (EGFR) IHC ($102,400\times 93,184$ pixels) stained tissue microarrays, and each comprised 105 tissue samples corresponding to 105 different patients. Totally 105 H&E stained tissue images with their consecutive EGFR counterparts were cropped from the microarrays and the image sizes vary from $7,000\times 7,000$ to $10,000\times 10,000$ pixels. We randomly selected 55 samples as the training set while the other 50 pairs as the testing set.

An extra natural image dataset was used to validate the generalizability of our methodology. We collected 17 and 20 high-resolution ($3456\times 5184$ pixels) unpaired images taken in Tokyo during summer and autumn, respectively, as the training set and additional four summer images were used as a testing set. This Kyoto summer2autumn dataset¹¹1Kyoto summer2autumn dataset is available at: https://github.com/Kaminyou/Kyoto-summer2autumn was released to facilitate solving ultra-high-resolution-related problems that most computer vision studies might encounter.

Three popular unpaired image-to-image translation frameworks: CycleGAN [33], CUT [28], and L-LSeSim [32], were utilized to verify our approach. We followed the hyperparameter settings described in the original papers during the training process except for the model output size, which was changed from $256\times 256$ to $512\times 512$. We trained the CycleGAN, CUT, and L-LSeSim for 50, 100, and 100 epochs. Models were trained and tested on three datasets: ANHIR, glioma, and Kyoto summer2autumn. Due to the insufficiency of WHIs in ANHIR dataset, we could only inference on the training set (note that training was in an unsupervised manner) while glioma and Kyoto summer2autumn datasets can be further split into training and testing sets. We replaced all IN layers with KIN layers in the generators during the inference process. One ultra-high-resolution image would be cropped into non-overlapped patches and pass through the KIN module. Translated patches were assembled to the final translated output. Constant and Gaussian kernels with sizes of 3, 7, and 11 were used to generate the best results. Translated images generated with KIN were compared with those from IN and TIN. Due to the GPU memory limitation, translated images generated with IN were also in a patch-wise (512$\times$512) manner, which is the same as the patch-wise IN in Chen et al. ’s work [8].

In addition to the visualization of the translated images, we calculated Fréchet inception distance (FID) [14], histogram correlation, Sobel gradients [16] in YCbCr color domain, perception image quality evaluator (PIQE) [30], and natural image quality evaluator (NIQE) [25] to comprehensively evaluate the quality of translated ultra-high-resolution images.

However, due to the limitations of the available metrics that the tiling artifacts are difficult to be fairly graded and the unavailability of the perfectly matched counterpart, we conducted two human evaluation studies with five specialists: (a) quality challenge: given one source, one reference, and three translated images generated by patch-wise IN, TIN, and our KIN methods, respectively, specialists were asked to select the best among three translated images in 30 seconds. Since the images generated by CycleGAN and L-LSeSim were atrocious, we only chose images generated by CUT; (b) fidelity challenge: given one real image and one translated image, specialists were asked to select the one which is personally considered realistic in 10 seconds. We followed the protocol of the AMT perceptual studies from Isola et al.  [17] but adjusted the time limitation as our images are extremely large. Since the data in ANHIR breast, COAD, and lung lesion subdatasets are insufficient, we combined these subdatasets as single ANHIR dataset and randomly selected pairs of real and translated WSIs from it.