Adversarial Stain Transfer to Study the Effect of Color Variation on Cell Instance Segmentation

The data used for the color transfer study are mouse neuron microscopic images. Histological sections of a mouse brain were stained using IHC: sections were incubated with NeuN antibody, then treated with Diaminobenzidine (DAB) chromogen and nickel (a DAB enhancer) to develop color. DAB staining is usually brown, and the addition of nickel helps the staining become darker and easier to recognize. The DAB amount stayed constant throughout the staining process, and the only variable was the nickel concentration. This procedure yielded five section groups with different color representations, one of which was nickel-free and the other four were enhanced with nickel concentrations of 0.4%, 0.8%, 1.2% and 1.6%, respectively. The nickel concentration in this study is estimated by the proportion of nickel solution to DAB solution. Fig. 1 left illustrates image examples from the five groups. The higher the concentration of nickel, the darker the color of neurons. We manually extracted 10 images (3k $\times$ 3k pixels) of the cortex from each section group to build our dataset, including a training set of 13k 256 $\times$ 256 images and a test set of five 512 $\times$ 512 images, which represented five color groups.

The cell segmentation dataset and a trained DL segmentation model were retrieved from [wu2022general]. Microscopic images were extracted from a macaque brain section stained with NeuN and nickel concentration of 1.6%, knowing that the distribution and staining of neurons in mouse and macaque cortex are similar. This dataset was independent of the mouse color transfer dataset. In this study, we used only a subset of cortex, containing 36 1024 $\times$ 1024 images, on which we altered the stain color to examine the effect on segmentation.

We converted images from RGB to log-chroma space and calculated the histograms in this space. The intensity of each channel was normalized by the other two channels to make the histogram differentiable and insensitive to illumination changes. For example, the conversion of the R channel is defined as [afifi2021histogan]:

where $I_K$ is the intensity of the $K$ channel of image $I$, $ϵ$ is added for numeric stability, and the $u$ and $v$ are the $K$ channel normalized by the other two channels, respectively. The $G$ and $B$ channels are also projected into log-chroma space to compute $I_{u}G$, $I_{v}G$, $I_{u}B$, $I_{v}B$. Binning is performed in u and v dimensions for each channel, resulting in a histogram with the shape of $n\times n\times 3$, where $n$ equals 64, is the number of bins.

The difference between the two histograms was measured using Hellinger distance (H dist) and Kullback-Leibler divergence (KL div) [afifi2021histogan]. They are usually used to quantify the similarity between two probability distributions. The more similar the distributions are, the closer the value is to 0. Previous studies suggested that H dist had a better sensibility to minor deviations, whereas KL div was more responsive to large deviations [afifi2021histogan, mocnik2021benford]. In this study, they were used to assess the color variation in different nickel groups and compare the color transfer performance of different normalization methods. In particular, the H dist was also employed to train the neural network as the color matching loss.

The objective of ReHistoGAN was to generate a realistic image, preserving both the content of the source image and the color information of the target image. As shown in Fig. 2, the network consisted of a generator to recolor the input image and a discriminator. The generator network was modified from a U-Net-like structure with skip connections. The target histogram features were projected into latent space and inserted in the last two blocks of the decoder to control the color of the output image. The generator was trained to generate an image having the same content as the source image and a color distribution similar to the target image. It is worth noting that the ReHistoGAN was trained in a fully unsupervised manner since the only ground truth needed was the histogram features, which were computed as described in 2.2. The goal of the discriminator was to distinguish the generated image from the real one. The loss function of the entire network is defined as [afifi2021histogan]:

where $t$, $s$ and $r$ are target, source and recolored images respectively. $H$ is the histogram. $L_C$ is the color matching loss computed with Hellinger distance. $I$ is perceptual detail (without color information), which was obtained with the Laplacian operator. $L_R$ is the reconstruction loss computed with the L1 norm. The discriminator loss $L_D$ is used to measure how realistic the generated image is. $\alpha$, $\beta$ and $\gamma$ are hyperparameters equal to 32, 1.5 and 4, respectively, which were defined empirically in [afifi2021histogan] to control the weight of each term.

The ReHistoGAN model was then applied to recolor each image of the cell segmentation dataset using itself and five images in the mouse color transfer test set as targets. This procedure yielded six supplementary groups with continuous color changes. The first was a control group where no color changes were expected (same source image as target image), which was added to measure possible artifacts brought by the GAN network during encoding and decoding. The remaining five groups inherited the color distributions from the color transfer dataset, which were different from the segmentation set. Six colors on the same set allowed us to study the effect of color variation on cell segmentation.

We utilized the segmentation method proposed in [wu2022general], which consisted of a neural network to predict the cell, the inter-cell contour and the background, followed by a post-processing scheme based on mathematical morphology to individualize each cell instance. Color augmentation techniques were applied during the training to improve the robustness of the model to color inconsistencies. After training, we applied the network on the segmentation test set as well as the six GAN-recolored sets to study the impact of color variation.

In this study, the color variation between images was assessed by the difference between color histograms (H dist and KL div). A good color transfer method should provide a recolored image with low H dist and KL div matching the target.

The cell segmentation was evaluated using three metrics: F1 score, Aggregated Jaccard Index (AJI) and Dice [12]. F1 score evaluated the cell instance segmentation, a segmented cell was considered a true positive when it overlapped the ground truth with an intersection over union greater than 0.5. AJI was a stricter metric to evaluate the instance segmentation since it penalized false segmentations. We also computed the Dice score to evaluate the semantic segmentation.

We compared ReHistoGAN with four state-of-the-art color transfer methods: Macenko et al. [macenko2009method], Reinhard et al. [reinhard2001color], Vahadane et al. [vahadane2016structure] and histogram matching [coltuc2006exact]. Each image in the mouse test set was recolored using all color groups as color references, resulting in 25 combinations in total. Fig. 3 shows the visual comparison of two extreme cases from the nickel-free group (brown) and the highest-nickel group (dark). Macenko et al. [macenko2009method] changed the stain color while it introduced artifacts, especially when recoloring the dark stain to brown. Reinhard et al. [reinhard2001color] was unable to retain all color information of the target image, with the stain color in the results still affected by the source images. Vahadane et al. [vahadane2016structure] achieved better results visually, however both Macenko et al. [macenko2009method] and Vahadane et al. [vahadane2016structure] suffered intensity issues, the recolored images were unnaturally too bright compared to the source and target images. The possible explanation for this might be that the IHC stain is a scatterer of light rather than an absorber. Therefore the Beer-Lambert law of absorption used for color space conversion is no longer applicable [vahadane2016structure]. The dark stain was successfully recolored into brown using the histogram matching [coltuc2006exact]. On the other hand, it did not perform well in converting brown stains to dark (brown spots still exist, see Fig. 3), and the resulting stain colors were absent from both the source and target images. Compared to other methods, ReHistoGAN showed better performance and robustness, and it worked well in all cases, transferring the target color without losing the source content.

Table 1 reports the quantitative comparison of ReHistoGAN and other methods on mouse data. H dist and KL div of color histograms were measured between the recolored images and their target images. Histogram matching had the best average H dist and KL div since the objective was to produce similar color histograms. However, the high standard deviation suggests that this method lacks robustness. This finding is consistent with visual results in Fig. 3, where the method performed poorly when the source and target images had very dissimilar statistics. On the other hand, besides comparable quantitative results as histogram matching (same KL div and slightly higher H dist), ReHistoGAN had the lowest standard deviations, which indicates its robustness. Despite bringing illumination changes, Macenko et al. [macenko2009method], Reinhard et al. [reinhard2001color] and Vahadane et al. [vahadane2016structure] were unable to correctly adjust color variation in IHC stained images. Compared to Histogram matching [coltuc2006exact] and ReHistoGAN, these methods resulted in recolored images with more significant color differences from the target images.

We applied ReHistoGAN to recolor images of the macaque segmentation test set, as illustrated in Fig. 4. The second column is the control group recolored using the same macaque image as the target, which allowed us to investigate the possible effect of encoding and decoding. The following columns show recolored images using five nickel mouse groups as the target, respectively. In total, six additional sets were produced, including one set without color change (control) and five with continuous color variations from brown to dark (recolored 1-5). Surprisingly, experts estimated that the original macaque images were closest in color to recolored-3 (0.8%), although the nickel concentration used for macaque images was the same as recolor-5 (1.6%). This inconsistency may be due to factors other than nickel concentration during the staining process.

We evaluated the segmentation network on the original and six additional sets of macaque images. Table 2 reports the segmentation results and the color variation between the test sets and the training set. Despite the fact that the model was trained using color augmentation techniques, segmentation performance has been negatively impacted by color variations. The scores of the original set were the best since they came from the same data as the training set. They were used as references to show the performance degradation due to GAN artifacts and color variation. The scores of the control group were comparable to the reference but slightly lower, indicating that almost all information of the original image was preserved during encoding and decoding. The group with the most degradation in segmentation was the recolored-1 set, which had the most significant color variation compared to the original set (KL div: 0.27). In particular, the AJI of this set decreased by 17% when compared to the reference. The recolored-3 set, on the other hand, scored similarly to the control group as it had similar color distributions both visually and quantitatively (KL div: 0.01). It suggests that the DL segmentation model is robust to slight color variation. Overall, the findings suggest that the segmentation performance decreases more in test sets with higher color variation. Surprisingly, the segmentation in the recolored-4 set was better than that in recolored-2, even though the former had a more visually distinct color variance. This might be due to the fact that despite the color variation, the contrast between cell and tissue is also essential for neural networks to segment cells. In our case, the darker stained images in the recolored-4 set presented a more significant contrast than lighter stained images in the recolored-2 set, providing a better condition for cell segmentation.