DiffMix: Diffusion Model-based Data Synthesis for Nuclei Segmentation and Classification in Imbalanced Pathology Image Datasets

SDM is a conditional denoising diffusion probabilistic model (CDPM) conditioned on semantic label maps. Based on CDPM, SDM follows two fundamental diffusion processes i.e., forward and reverse process. The reverse process is a Markov chain with Gaussian transitions. When the added noise is large enough, the reverse process is approximated by a random variable $\textbf{y}_{T}\sim{\mathcal{N}(0,\textbf{I})}$, defined as follows:

The forward process implements Gaussian noise addition for $T$ timesteps based on variance schedule $\{$${\beta_{1},...\beta_{T}}$$\}$ as below:

With $\alpha_{t}:=1-\beta_{t}$ and $\bar{\alpha_{t}}:=\prod_{s=1}^{t}\alpha_{s}$, we can write the marginal distribution as follows,

The conditional DDPM is optimized to minimize the negative log-likelihood of the data for the given input and condition information. If noise in the data follows Gaussian distribution with a diagonal covariance matrix $\Sigma_{\theta}(\textit{y}_{t},\textit{x},t)=\sigma_{t}\textbf{I}$, denoising can be the optimization target by removing the noise assumed to be present in data as follows,

SDM is a U-Net-based network that estimates the noise from the noisy input image. Unlike other conditional DDPMs, the denoising network of SDM processes the semantic label map $x$ and noisy input $y_{t}$ independently. While $y_{t}$ is fed into the encoder part, $x$ is injected into the decoder to fully leverage the semantic information[19]. As for training, SDM is trained in a manner similar to the improved DDPM [14] so that it not only predicts the involved noise to reconstruct the input image but also predicts variances to enhance the log-likelihood of the generated images. To improve sample quality, SDM utilized classifier-free guidance for inference.

SDM replaces the semantic label map $x$ with an empty (null) map $\emptyset$ in order to separate the noise estimated under the label map guidance by $\epsilon_{\theta}(y_{t}|x)$, from the noise estimated in an unconditioned case $\epsilon_{\theta}(y_{t}|\emptyset)$. The strategy allows for the inference of the gradient of the log probability, expressed as follows,

In the sampling process, the disentangled component $s$ is increased to improve the samples from conditional diffusion models, formulated as follows,

Fig 1 illustrates the process of creating custom label maps to condition the semantic diffusion model for synthesizing desired data based on the original input image label $y_{0}$. We have prepared custom semantic label maps to condition SDM to direct it to synthesize the data for our imbalanced datasets. So, we have considered making two types of semantic label maps to improve imbalanced datasets. First, balancing maps to balance the number of nuclei among different nuclei types. GradMix [5] have increased the fewest type nuclei in datasets by cutting, pasting, and smoothing for both images and labels. On the other hand, we only used their mixed labels for our experiment. Second, enlarging maps to stretch the data distribution available in training. We randomly moved nuclei positions on semantic maps to synthesize diverse image patches with SDM by conditioning with the unfamiliar semantic maps to lead the diffusion model to generate as many various patches as possible.

We synthesize the virtual data with pretrained SDM conditioned on the custom semantic label maps $x$. Fig 1 depicted the data sampling process in SDM. Before putting the original image $y_{0}$ into diffusion net $f$, we added noise to $y_{0}$. The semantic label and noisy image $y_{t}$ will be simultaneously used. To synthesize virtual images, we built two label maps and trained a semantic diffusion model. Before we start data synthesis, we add noise on the input image $y_{0}$. After that, we input $y_{t}$ and $x$ to the pretrained denoising network $f$, $y_{t}$ for the encoder, and $x$ for the decoder. As SDM generates samples, it uses the empty label $\emptyset$ to generate unconditioned output. The image is sampled from an existing but noised patch, depending on the pre-defined time steps. By doing so, we generate patches that are conditioned on the custom semantic label maps. In this process, we added noise on the input image $y_{0}$, so we input the noised input $y_{t}$ into the pretrained denoising network $f$, and following [19], we input the custom semantic label maps to decoder parts. Then the semantic label maps will condition SDM to synthesize image data that satisfies the label maps. We have sampled the new data with DDIM[17] process, which diminishes the sampling steps but with high-quality data.

In this study, we have used two imbalanced nuclei segmentation and classification datasets for our experiment. First, GLySAC [6] consists of 59 H&E images of size 1000$\times$1000 pixels, and split into 34 train images and 25 test images. The GLySAC has 30875 nuclei, and grouped into 3 nuclei types which are 12081 lymphocytes, 12287 epithelial and 6507 miscellaneous, respectively. Second, CoNSeP [8] consists of 41 H&E images of size 1000$\times$1000 pixels, and divided into 27 train images and 14 test images. The CoNSeP has 24319 nuclei in total, and composed of four nuclei classes, which are 5537 epithelial nuclei, 3941 inflammatory, 5700 spindle, and 371 miscellaneous nuclei.

We used one NVIDIA RTX A6000 to train SDM, we trained SDM for 10000epochs. For data synthesis, we implemented DDIM-based diffusion process from 1000 to 100, and we added noise on the input image to SDM, setting $T$ as 55. In our scheme, $\emptyset$ is defined as the all-zero vector as same as [19] and set $s=1.5$ when sampling both datasets. We implemented experiments on two baseline networks SONNET[6] and HoVer-Net[8]. SONNET is implemented with Tensorflow version 1.15[1] as software framework with two NVIDIA GeForce 2080 Ti GPUs. HoVer-Net is trained with PyTorch 1.11.0 as software framework with one NVIDIA GeForce 3090 Ti GPU. We implemented 4-fold cross validation for SONNET, and 5-fold cross validation for HoVer-Net. For fair comparison, we trained each process with same iterations, changing only the epoch numbers depending on the training set size. DiffMix and GradMix used all the original patches and the same numbers of synthesized patches In case of DiffMix-B, it has original data and balancing map based patches. Likewise, DiffMix-E training set consists of enlarging patches with original training set. Therefore, We trained each baseline network for 100epochs, 75epochs for DiffMix-B and -E, and trained 50epochs for GradMix and DiffMix.

Fig 2 presents a qualitative comparison of synthesized patches. The original patch is on the left, followed by enlarging, balancing, and GradMix patches. Our two types of patches are well-harmonized with the surrounding structure, compared to GradMix. Moreover, using our scheme, we can synthesize many patches using a semantic diffusion model.

For the quantitative evaluation, we implemented two state-of-the-art networks, HoVer-Net and SONNET, on two public imbalanced nuclei-type datasets, GLySAC and CoNSeP. Table 1 shows the results of 5 experiments per network for each dataset. We also conducted ablation studies on balancing (DiffMix-B) and enlarging (DiffMix-E) patch datasets. Before analyzing the experiment results, we computed the proportion of the least presenting nuclei type in each dataset. We found that Miscellaneous nuclei accounted for around 19$\%$ in GLySAC, but only 2.4$\%$ in CoNSeP. This means that GLySAC is more balanced in terms of nuclei types than CoNSeP. Taking this information into account, we analyzed our experiment results. First, we noticed that DiffMix-E showed the highest classification performance in the GLySAC dataset. This result indicates that enlarging semantic map-based data synthesis, like DiffMix-E, had enough opportunities to enlarge its learning distribution for GLySAC. However, in the CoNSeP dataset, DiffMix-E performed lower than other methods, suggesting that if a dataset is somewhat balanced, it is important to enlarge the available data distribution. DiffMix showed the highest performance in most metrics, with a 4$\%$ and 9$\%$ margin from the second-highest result in classifying Miscellaneous, successfully diminishing the classification performance variability among class types. Furthermore, DiffMix improved the segmentation and classification performance of two state-of-the-art networks, even compared to GradMix.