DARC: Distribution-Aware Re-Coloring Model for Generalizable Nucleus Segmentation

In this paper, we adopt a U-Net-based model similar to that in [1] as the baseline. It performs both semantic segmentation and contour detection for nucleus instances. The area of each nucleus instance is obtained via subtraction between the segmentation and contour prediction maps [1]. Details of the baseline model is provided in the supplementary material. To handle domain variations, we adopt IN rather than batch normalization (BN) in the U-Net model.

Our proposed Distribution-Aware Re-Coloring model (DARC) is illustrated in Fig. 2. Compared with the baseline, DARC replaces the IN layers with the proposed Distribution-Aware Instance Normalization (DAIN) layers. DARC first re-colors each image to relieve the influence caused by image acquisition conditions. The re-colored image is then fed into the U-Net encoder and the ratio prediction head. This head predicts the ratio between foreground and background pixel numbers. With the predicted ratio, the DAIN layers can estimate feature statistics more robustly and facilitate more accurate nucleus segmentation.

We propose the Re-Coloring (RC) method to overcome the color change in different domains. Specifically, given a RGB image $\boldsymbol{I}$, e.g., an H$\&$E or IHC stained image, we first obtain its grayscale image $\boldsymbol{I}_{g}$. We then feed $\boldsymbol{I}_{g}$ into a simple module $T$ that consists of a single residual block and a $1\times 1$ convolutional layer with output channel number of 3. In this way, we obtain an initial re-colored image $\boldsymbol{I}_{r}$.

However, de-colorization results in the loss of fine-grained textures and may harm the segmentation accuracy. To handle this problem, we compensate $\boldsymbol{I}_{r}$ with the original semantic information contained in $\boldsymbol{I}$. Recent works [40] show that semantic information can be reflected via the order of pixels according to their gray value. Therefore, we adopt the Sort-Matching algorithm [41] to combine the semantic information in $\boldsymbol{I}$ with the color values in $\boldsymbol{I}_{r}$. Details of RC is presented in Alg. 1, in which $Sort$ and $ArgSort$ denote channel-wisely sorting the values and obtaining the sorted values and indices respectively, and $AssignValue$ denotes re-assembling the sorted values according to the provided indices. Details of the module $T$ are included in the supplementary material.

Via RC, the original fine-grained structure information from $\boldsymbol{I}_{g}$ is recovered in $\boldsymbol{I}_{r}$. In this way, the re-colored image is advantageous in two aspects. First, the appearance difference between pathological images caused by the change in scanners and staining protocols is eliminated. Second, the re-colored image preserves fine-grained structure information, enabling precise instance segmentation to be possible.

Due to dramatic domain gaps, feature statistics may differ significantly between domains [7, 5, 6], which means that feature statistics obtained from the source domain may not apply to the target domain. Therefore, existing DG works usually replace BN with IN for feature normalization [19, 12]. However, for dense-prediction tasks like semantic segmentation or contour detection, adopting IN alone cannot fully address the feature statistics variation problem. This is because feature statistics are also relevant to the ratio between foreground and background pixel numbers. Specifically, an image with more nucleus instances produces more responses in feature maps and thus higher feature statistic values, and vice versa. The difference in this ratio causes interference to nucleus segmentation.

To verify the above viewpoint, we evaluate the baseline model under different foreground-background ratios. Specifically, we first remove the foreground pixels via in-painting [27], and then pad the original testing images with the obtained background patches. We adopt $B$ to denote the ratio between the size of the obtained new image and the original image size. Compared with the original images, the new images have the same foreground regions but more background pixels, and thus have different foreground-background ratios. Finally, we evaluate the performance of the baseline model with different $B$ values. Experimental results are presented in Table 1. It is shown that the value of $B$ affects the model performance significantly.

The above problem is common in nucleus segmentation because pathological images from different organs or tissues tend to have significantly different foreground-background ratios. However, this phenomenon is often ignored in existing research. To handle this problem, we propose the Distribution-Aware Instance Normalization (DAIN) method to re-estimate feature statistics that account for different ratios of foreground and background pixels. Details of DAIN is presented in Alg. 2. The structures of $E_{\boldsymbol{\mu}}$ and $E_{\boldsymbol{\delta}}$ are included in the supplemental materials.

As shown in Fig. 2, to obtain the foreground-background ratio $\boldsymbol{\rho}$ of one input image, we first feed it to the model encoder with $\boldsymbol{\Delta s_{ra}}$ as the additional input. $\boldsymbol{\Delta s_{ra}}$ acts as pseudo residuals of feature statistics and is obtained in the training stage via averaging $\boldsymbol{\Delta s}$ in a momentum fashion. The output features by the encoder are used to predict the foreground-background ratio $\boldsymbol{\rho}$ with a Ratio-Prediction Head (RPH). $\boldsymbol{\rho}$ is then utilized to estimate the residuals of feature statistics: $\boldsymbol{\Delta s}=f(\boldsymbol{\rho})$. Here, $f$ is a $1\times 1$ convolutional layer that transforms $\boldsymbol{\rho}$ to a feature vector whose dimension is the same as the target layer’s channel number. After that, the input image is fed into the model again with $\boldsymbol{\Delta s}$ as additional input and finally makes more accurate predictions.

The training of RPH requires an extra loss term $L_{rph}$, which is formulated as bellow:

where $\boldsymbol{\rho_{g}}$ denotes the ground truth foreground-background ratio, and $L_{BCE}$ and $L_{MSE}$ denote the binary cross entropy loss and the mean squared error, respectively.

The proposed method is evaluated on four datasets, including two H$\&$E stained image datasets CoNSeP [3] and CPM17 [28] and two IHC stained datasets DeepLIIF [29] and BC-DeepLIIF [29, 32]. CoNSeP [3] contains 28 training and 14 validation images, whose sizes are $1000\times 1000$ pixels. The images are extracted from 16 colorectal adenocarcinoma WSIs, each of which belongs to an individual patient, and scanned with an Omnyx VL120 scanner within the department of pathology at University Hospitals Coventry and Warwickshire, UK. CPM17 [28] contains 32 training and 32 validation images, whose sizes are $500\times 500$ pixels. The images are selected from a set of Glioblastoma Multiforme, Lower Grade Glioma, Head and Neck Squamous Cell Carcinoma, and non-small cell lung cancer whole slide tissue images. DeepLIIF [29] contains 575 training and 91 validation images, whose sizes are $512\times 512$ pixels. The images are extracted from the slides of lung and bladder tissues. BC-DeepLIIF [29, 32] contains 385 training and 66 validation Ki67 stained images of breast carcinoma, whose sizes are $512\times 512$ pixels.

In the training stage, patches of size $224\times 224$ pixels are randomly cropped from the original samples. During training, the batch size is 4 and the total number of training iterations is 40,000. We use Adam algorithm for optimization, and the learning rate is initialized as $1e^{-3}$, which is gradually decreased to $1e^{-5}$ during training. We adopt the standard augmentation, like image color jittering and Gaussian blurring. In all experiments, the segmentation and contour detection predictions are penalized using the binary cross entropy loss.

In this paper, the models are compared using the AJI [33] and Dice scores. In the experiments, models trained on one of the datasets will be evaluated on the three unseen ones. To avoid the influence of the different sample numbers of the datasets, we calculate the average scores within each unseen domain respectively and then average them across domains.

In this paper, we re-implement some existing popular domain generalization algorithms for comparisons under the same training conditions. Specifically, we re-implement the TENT [34], BIN [19], DSU [20], Frequency Amplitude Normalization (AmpNorm) [36, 37], SAN [35] and EFDMix [40]. We also evaluate the stain normalization [38] and stain mix-up [39] methods that are popular in pathological image analysis. Their performances are presented in Table 2. DARC_all replaces all normalization layers with DAIN, while DARC_enc replaces the normalization layers in the encoder with DAIN and uses BN in its decoder. As shown in Table 2, DARC_enc achieves the best average performance among all methods. Specifically, DARC_enc improves the baseline model’s average AJI and Dice scores by $4.81\%$ and $7.04\%$. Compared with the other domain generalization methods, DAIN, DAIN w/o Ratio, DARC_all and DARC_enc achieve impressive performances on BC-DeepLIIF, which justify that re-estimating the instance-wise statistics is important for improving the domain generalization ability of models trained on BC-DeepLIIF. Qualitative comparisons are presented in Fig. 3. Moreover, the complexity analysis between the baseline model and DARC_enc is presented in Table 3.

We separately evaluate the effectiveness of RC and DAIN, and present the results in Table 2. Also, we train a variant model without foreground-background ratio prediction, which is denoted as ‘DAIN w/o Ratio’ in Table 2. Compared with the baseline model, RC improves the average AJI and Dice scores by 1.41$\%$ and 2.59$\%$, and DAIN improves the average AJI and Dice scores by 1.13$\%$ and 4.08$\%$. Compared with the variant model without foreground-background ratio prediction, DAIN improves the average AJI and Dice scores by 0.90$\%$ and 3.74$\%$. Finally, the combinations of RC and DAIN, i.e., DARC_all and DARC_enc, achieve the best average scores. As shown in Table 2, DARC_enc improves DARC_all by 1.24$\%$ and 0.77$\%$ on AJI and Dice scores respectively. This is because after the operations by RC and DAIN in the encoder, the obtained feature maps are much more robust to the domain gaps, which enables the decoder to adopt the fixed statistics maintained during training. Moreover, using the fixed statistics is helpful to prevent the decoder from the influence of varied foreground-background ratios on feature statistics.