Prior-guided Diffusion Model for Cell Segmentation in Quantitative Phase Imaging

Diffusion models (DM) are advanced generative models capable of generating high-quality images [12, 22] and effectively learning useful semantic information [23, 24]. The training process of a DM involves two main steps: forward diffusion and reverse diffusion. In the forward diffusion process, a transition is incrementally made from the original data distribution towards a Gaussian distribution across a sequence of discrete time steps $t$, which range from $0$ to $T-1$. At each step, a predetermined amount of Gaussian noise $\bm{\epsilon}$ is added to the input images $\bm{x}$, progressively transforming the data as follows:

where $\bar{\alpha}_{t}$ is a hyper-parameter, $\{\bm{x_{t}}\}^{T}_{t=1}$ are noise-added images and $T$ is the maximal sampling steps. During the reverse diffusion process, the denoising U-Net denoted as $\bm{\epsilon_{\theta}}$ is responsible for predicting the noise in $\bm{x_{t}}$. The training objective is to minimize the error between the added noise $\bm{\epsilon}$ and the predicted noise $\bm{\epsilon_{\theta}}\left(\bm{x_{t}}\right)$, which can be represented as follows:

Implementing a traditional DM poses challenges due to high computational costs. To address this challenge, the latent diffusion model (LDM) [13] was proposed. This method incorporates the idea of compression by introducing pre-trained autoencoding models to encode images into a latent space. Therefore, the training of the diffusion model in the LDM method is more efficient. Specifically, a pre-trained autoencoding model includes an encoder $\bm{\mathcal{E}}$ and a decoder $\bm{\mathcal{D}}$. All input images $\bm{x}$ for LDM are encoded into the latent representation $\bm{z}$ by $\bm{\mathcal{E}}$. The following diffusion training process is in the latent space. The training objective can be represented as follows:

Stochastic DDPM sampling [12] and deterministic DDIM sampling [19] are two commonly used sampling methods for a trained DM to generate segmentation masks. Specifically, DDPM sampling follows a Markov chain process, and the sampling process is as follows.

where $\bm{\epsilon_{\theta}}$ represents pre-trained denoising U-Net within the DM, $\bm{\epsilon}$ represents random sample Gaussian noise. The hyper-parameters $\alpha_{t}$, $\bar{\alpha}_{t}$, $\beta_{t}$ and $\sigma_{t}$ are used to determine the scales of noise. Due to the inclusion of random noise $\bm{\epsilon}$, the sampling process in DDPM is inherently stochastic.

In contrast, DDIM sampling is a deterministic process, the formulation is as follows:

where $\bm{\epsilon_{\theta}}$ represents pre-trained denoising U-Net within the DM, while the hyper-parameter $\bar{\alpha}_{t}$ determines the scales of the noise. The DDIM process does not involve random noise $\bm{\epsilon}$. This process is defined as a deterministic sampling method and can be reformulated to derive its corresponding ordinary differential equation (ODE) [19]. Consequently, numerical approximation methods, including the Euler method, are capable of approximating the inversion process of DDIM sampling. By using the Euler method with a sufficient number of discrete time steps, the DDIM inversion is effectively deterministic.

The deterministic characteristic of DDIM sampling and DDIM inverse sampling ensures consistent transformation between input images and their corresponding noise. To facilitate the recovery of the original image from noise, the DDIM inverse sampling method preserves important content information during its noise addition process. Inspired by the ability to retain content information, DDIM inversion is employed in areas such as image translation [25, 26] and image editing [27] for extracting content information of input images.

As shown in Fig. 1, the proposed method is named the prior-guided DM-based segmentation (PG-DiffSeg). Leveraging the high-contrast content information in QPI images, this method incorporates content information into the starting noise of the sampling process, enabling DM-based segmentation to achieve precise results without multiple samplings and ensemble learning. Specifically, this method utilizes two latent diffusion models (LDMs) [13]: one LDM for prior extraction (LDM-P) and another for segmentation (LDM-S). By exploiting the deterministic DDIM inversion process, LDM-P transforms to-be-segmented images into corresponding noise, which conforms to a Gaussian distribution while retaining content information. Instead of relying on randomly sampled noise, LDM-S utilizes the prior-informed noise extracted by LDM-P as the starting noise to generate the segmentation mask. Starting from the prior-informed noise better supports LDM-S in stably generating a precise segmentation mask.

To generate Gaussian noise while retaining content information in images, a latent diffusion model (LDM) with the DDIM inverse sampling process is employed to transform input images into their corresponding noises. The LDM used in this process is named LDM-P. Specifically, the training of LDM-P involves both the forward and reverse diffusion processes in an unsupervised manner. For efficient training, the input images are transformed into latent representations $\bm{z}$ using a pre-trained encoder $\bm{\mathcal{E}}$ [13]. During the forward diffusion process, random noise is introduced to $\bm{z}$, as illustrated in Eq.1. Subsequently, during the reverse diffusion process, LDM-P predicts the added noise by minimizing the training loss defined in Eq. 3. Given to-be-segmented images, the trained LDM-P model converts the input latent representation into prior-informed noise for DM-based segmentation described below.

Given the prior-informed noise generated by LDM-P, another LDM is employed to predict the segmentation mask of the to-be-segmented image. In our study, this LDM for segmentation tasks is referred to as LDM-S. Training LDM-S involves both images and segmentation masks and follows the supervised learning process. The paired inputs are first pre-processed. The to-be-segmented images $\bm{x}$ are transferred to latent representation $\bm{z}$ using the pre-trained encoder $\bm{\mathcal{E}}$ [13]. The segmentation mask $\bm{y}$ is initially converted into multiple binary segmentation masks through one-hot encoding and then resized to match the dimensions of the latent representation $\bm{z}$. The training of LDM-S involves both forward and reverse diffusion processes. In the forward diffusion, Gaussian noise is added to the pre-processed segmentation mask $\bm{y}$, as illustrated in Eq. 1. During the reverse diffusion, with the condition of the to-be-segmented image, a denoising U-Net $\bm{\epsilon_{\theta}}$ is utilized to predict the added noise on the segmentation masks $\bm{y_{t}}$. The training objective aims to minimize the loss function, defined as the squared error between the added noise $\bm{\epsilon}$ and the noise predicted by the model, $\bm{\epsilon_{\theta}}(\bm{y_{t}},\bm{z})$:

The trained LDM-S model predicts the segmentation mask using the prior-informed noise $\bm{\epsilon_{p}}$ as the starting noise, which is generated by the trained LDM-P on the given to-be-segmented image. The segmentation in LDM-S requires the dimensions of the starting noise to align with those of the segmentation mask. The prior-informed noise $\bm{\epsilon_{p}}$ obtained from LDM-P has a dimension of $\bm{\epsilon_{p}}\in\mathbb{R}^{h/f\times w/f\times 3}$, where $h\times w$ corresponds to the original image size, and $f$ signifies the down-sampling factor. Averaging and replication operations are performed on $\bm{\epsilon_{p}}$ to preprocess it into $\bm{\epsilon_{p}}\in\mathbb{R}^{h/f\times w/f\times n}$, where $n$ represents the number of classes in the segmentation mask.

Using this pre-processed starting noise, the trained LDM-S employs the DDIM sampling process (Eq. 5) to generate segmentation mask predictions. Subsequently, a resizing operation is applied to convert the prediction from the latent space size to the image space. The testing process of LDM-S is presented in Algorithm 1.

Two QPI datasets for cell viability assessment study [6], HeLa and CHO-2, were utilized to evaluate model performance on semantic segmentation, aiming to distinguish between live and dead cells. For simplicity, CHO-2 was referred to as CHO in this study. The image size is 832$\times$832 pixels, each accompanied by a corresponding segmentation mask. The HeLa dataset contains 1,199 images, while the CHO dataset comprises 2,051 images. The original images were subdivided into smaller sizes of 256$\times$256 and split into training, validation and testing datasets. Details about the data can be found in Table 1.

The LDM-P and LDM-S were used for prior extraction and segmentation, respectively. The denoising U-Net within both LDMs follows the architecture designed in LDM [13], with the same number of U-Net layers, ResNet-structured blocks, and self-attention blocks. The model architecture for comparative methods is described in Sec. 4.4. The LDM-P prioritizes learning overall content information with a larger receptive field. Accordingly, it features larger hidden layers with dimensions of 224, 448, 672, and 896, along with higher down-sampling factors where the spatial resolution of input features to attention modules is set to 8, 16, and 32. Conversely, the LDM-S emphasizes learning structure and local information, necessitating larger spatial resolution in self-attention blocks. Accordingly, the spatial resolution of input features to attention modules is set to 16, 32, and 64. To reduce the computational burden, it incorporates a smaller hidden layer with dimensions of 128, 256, 384, and 512. The parameters for LDM-P and LDM-S are 274M and 84.6M, respectively.

In the experiments, the parameters predefined by LDM [13] were utilized. The configuration involved a maximum diffusion step of 1000, and the pre-trained weight named “vq-f4” was selected for the encoder $\bm{\mathcal{E}}$ to transform the size of input images from 256 to 64. In the training process, the LDM-P used a batch size of 20 and a learning rate of 2e-06, while the LDM-S utilized a batch size of 10 and a learning rate of 2e-06. The training for both LDMs involved standard data augmentation techniques, including rotation, horizontal and vertical flipping, as well as random adjustments to brightness and contrast. During testing, the LDM-P utilized a 100-step DDIM inversion to transform input images into prior-informed noise. The LDM-S employed a 200-step DDIM sampling (Eq. 5) in generating the segmentation mask.

Five state-of-the-art segmentation methods were implemented for comparison. Three of them are deterministic segmentation methods, while the other two are diffusion model-based segmentation methods using multiple sampling and ensemble learning.

Mean Intersection over Union (mIoU) and F1 score (F1) were used to evaluate the segmentation performance. Higher values for these metrics indicate better performance. The Structural Similarity Index (SSIM) was used to assess the content information retained in the starting noise for sampling. SSIM takes into account structure, texture, and luminance information, ranging from -1 to 1, with 1 indicating perfect similarity. Additionally, the Probability Density Function (PDF) of the starting noise is estimated using the Gaussian Kernel Density Estimation technique. Kullback-Leibler Divergence (KLD) is applied to determine the difference between the estimated PDF of the starting noise and that of the standard Gaussian distribution. The value of KLD ranges from 0 to positive infinity, with 0 indicating a standard Gaussian distribution. Both content information assessment and distribution assessment aim to analyze the superiority of the prior-informed noise in retaining content image information while still conforming to a Gaussian distribution. With three-fold cross-validation, all evaluation metrics included the mean and standard deviation (SD). Independent samples t-tests were conducted to identify significant differences in the results and were reported with p-values.

Experiments were conducted on two label-free imaging datasets to assess the effectiveness of the method in stratifying and contouring live and dead cells. Section 4.1 provided the dataset details, while Sections 4.4 and 4.5 discussed the comparative methods and evaluation metrics, respectively. The results are summarized in Table 2. It is worth noting that “Random$\times 1$” represents DM-based segmentation with only one sampling, while “Random$\times$3/$\times$5” and “CCDM$\times$3/$\times$5” represent 3 or 5 times of multiple sampling. Majority voting is employed to obtain ensemble predictions from these multiple predictions. Our proposed PG-DiffSeg is denoted as “Prior”.

1) Compared to two U-Net-based segmentation methods, the DM-based method with single sampling (“Random$\times$1”) can achieve comparable performance. By using multiple samplings in the DM-based method, the segmentation performance increased compared to “Random$\times$1” and also achieved better results than all the deterministic segmentation methods. This observation confirms that multiple sampling plays an important role in improving the performance of DM-based segmentation methods.

2) Nonetheless, increasing the number of multiple sampling does not always guarantee performance improvement. For example, from “Random$\times$1” to “Random$\times$3”, the mIoU and F1 increased 3.5% and 3.1% in HeLa dataset, respectively. However, from “Random$\times$3” to “Random$\times$5”, the mIoU and F1 remained the same performance. Same trend was also observed on CCDM. The possible reason is after a certain point, additional samplings may not provide significant new information or may even introduce noise into the segmentation process, resulting in a negative effect in performance. Compared to “Random”, the “CCDM” method showed higher performance on the HeLa dataset but not the CHO dataset. The distinction between ”Random” and ”CCDM” lies in the use of a pre-trained DINO model with a large amount of natural images to guide the diffusion training process in CCDM. The content information extracted by DINO had a positive impact on the HeLa dataset, whereas it had a negative effect on the CHO dataset. This difference may be attributed to the more complex structure of cell images within the CHO dataset, which significantly differs from the characteristics of the natural images used for pre-training DINO, leading to the ineffectiveness of the CCDM method.

3) Our proposed prior-guided method for DM-based segmentation significantly enhanced performance in terms of mIoU and F1 over both the HeLa and CHO datasets. Furthermore, by incorporating content information specific to the to-be-segmented images into the starting noise, only a single sampling was required to achieve optimal performance.

Visualization was used to demonstrate the effectiveness of the proposed method. The results are shown in Fig. 2. Additionally, the visualization of content prior is examined in Fig. 3.

In Fig. 2, the HeLa and CHO datasets were obtained using label-free imaging, focusing on distinguishing live cells, dead cells, and the background through semantic segmentation. It is clear from the dashed box that adding random noise to DM-based segmentation generates multiple different predictions. Using ensemble learning, like the majority vote, can improve overall prediction performance. However, these multiple predictions faced the problem of unpredictable variance, potentially misleading ensemble learning. Our method accurately predicted results using prior information with just one sample. Compared to deterministic methods such as U-Net, our approach also achieved better segmentation accuracy.

In Fig. 3, a comparison among random sampling (Gaussian distribution), the non-trainable method (forward diffusion), and the trainable method (DDIM inversion) is presented. The canny edge detector [34] is employed to analyze content information, while contrast enhancement followed by Gaussian blur is used to preprocess the starting noise. Both the HeLa and CHO datasets are utilized for content information and distribution analysis. From the histogram analyses, it can be observed that both the starting noise obtained by non-trainable and trainable methods follow a Gaussian distribution similar to Gaussian noise. Additionally, the contour information is retained in the starting noise for these two methods, while Gaussian noise does not retain any content information due to the randomness. Moreover, the red box indicates that the trainable method can retain more contour information than the non-trainable method, which is significantly important for segmentation.

The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.

Code and data will be made publicly available upon acceptance of the paper.