ArSDM: Colonoscopy Images Synthesis with Adaptive Refinement Semantic Diffusion Models

Unlike the previous generative methods, our work aims to generate a synthetic image with an identical segmentation mask to the original annotation. To accomplish this, we adapt the widely used conditional U-Net architecture [21] in the reverse process, where the mask is fed as a condition. Specifically, for an input image $\mathbf{x}_{0}\in\mathbb{R}^{H\times W\times C}$, $\mathbf{x}_{t}$ can be sampled at any timestep $t$ with the closed form:

where $\bm{\epsilon}\sim\mathcal{N}\left(\mathbf{0},\mathbf{I}\right),\alpha_{t}:=1-\beta_{t}$ and $\bar{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}$. It will be fed into the encoder $\mathcal{E}$ of the U-Net, and its corresponding mask annotation $\mathbf{c}_{0}\in\mathbb{R}^{H\times W}$ will be injected into the decoder $\mathcal{D}$. The model output can be formulated as:

Thus, the U-Net model $\bm{\epsilon}_{\theta}$ in Eq. 3 becomes $\bm{\epsilon}_{\theta}\left(\mathbf{x}_{t},t,\mathbf{c}_{0}\right)$, and the loss function in Eq. 3 is changed to:

The polyp regions in the colonoscopy images differ from the background regions, which contain more pathological information and should be adequately treated to learn a better model. However, training the diffusion models using the original loss function ignores the difference between different regions, where each pixel shares the same weights when calculating the loss. This would lead to the model generating more background-like polyps since the large background region will easily overwhelm the small foreground polyp regions during training. A simple way to alleviate this problem is to apply a weighted loss function that assigns the polyp and background regions with different weights. However, most polyps vary a lot in size and shape. Thus assigning constant weights for all polyps exacerbated the imbalance problem. In this case, to tackle this problem, we propose an adaptive loss function that vests different weights according to the size ratio of the polyp over the background. Specifically, we define a pixel-wise weights matrix $W^{\lambda}\in\mathbb{R}^{H\times W}$ with each entry $w^{\lambda}_{i,j}$ to be:

where $p=1$ means the pixel $p$ at $(h,w)$ belongs to the polyp region and $p=0$ means it belongs to the background region. Thus, the loss function becomes:

The downstream tasks of polyp segmentation and detection require rich semantic information on polyp regions to train a good model. Through extensive experiments, we found inaccurate sample images with coarse polyp boundary that is not aligned properly with the original masks may introduce large biases and noises to the datasets. The model can be confused by several conflicting training images with the same annotation. To this end, we design a refinement strategy that uses the prediction of a pre-trained segmentation model on the sampled images to guide the training process and restore the proper polyp boundary information. Specifically, at each iteration of training, the output $\tilde{\bm{\epsilon}}=\bm{\epsilon}_{\theta}\left(\mathbf{x}_{t},t,\mathbf{c}_{0}\right)$ will go into the sampler to generate sample image $\tilde{\mathbf{x}}_{0}$. Then, we take the sample image as the input of the segmentation model to predict the pseudo masks $\tilde{\mathbf{c}}_{0}$. We propose the following refinement loss based on IoU loss and binary cross entropy (BCE) loss between $\tilde{\mathbf{c}}_{0}$ and $\mathbf{c}_{0}$. The refinement loss is:

where $\mathcal{L}=\mathcal{L}_{IoU}+\mathcal{L}_{BCE}$ is the sum of the IoU loss and BCE loss, $\tilde{\mathbf{c}}_{0}$ is the collection of the three side-outputs ($\tilde{\mathbf{c}_{3}},\tilde{\mathbf{c}_{4}},\tilde{\mathbf{c}_{5}}$) and the global map $\tilde{\mathbf{c}_{g}}$ as described in [5]. $\mathcal{P}(\cdot)$ represents the PraNet model and $\mathcal{S}(\cdot)$ is the DDIM [16] sampler. The detailed procedure of one training iteration is shown in Algorithm 1 and the overall loss function is defined as:

We conducted our experiments on five public polyp segmentation datasets: EndoScene [20], CVC-ClincDB/CVC-612 [1], CVC-ColonDB [18], ETIS [14] and Kvasir [7]. Following the standard of PraNet, 1,450 image-mask pairs from Kvasir and CVC-ClinicDB are taken as the training set. The evaluations are conducted on the five datasets separately to verify the learning and generalization capability. The training image-mask pairs are padded to have the same height and width and then resized to the size of $384\times 384$. Experiments with prediction-guided sample refinement are trained with around one-half NVIDIA A100 days, while others are trained with approximately one day for convergence. We use the DDIM sampler with a maximum timestep of 200 for sampling images.

We conduct the evaluation of our methods and the state-of-the-art counterparts on polyp segmentation and detection tasks. We generated the same number of samples as the diffusion training set using the original masks, and then combined them to create a new downstream training set. We employed PraNet [5], SANet [22], and Polyp-PVT [2] as baseline segmentation models with default settings, and evaluated them using mean Intersection over Union (IoU) and mean Dice metrics. For detection, we selected CenterNet [30], Sparse-RCNN [17], and Deformable-DETR [31] as baseline models with the same settings as the original papers, and evaluated them using Average Precision (AP) and F1-scores.

The experimental results presented in Table 1 and 2 demonstrate the effectiveness of our proposed method in training better downstream models to achieve superior performance. Specifically, data generated by our approach assists the significant improvements for each model in mDice and mIoU, with increases of 6.0% and 5.7% over PraNet, 2.1% and 2.7% over SANet, and 0.7% and 0.7% over Polyp-PVT. We also observe superior AP and F1-scores compared to CenterNet, Sparse-RCNN, and Deformable-DETR trained with original data, with gains of 9.1% and 5.3%, 2.7% and 5.8%, and 3.4% and 6.1%, respectively. Moreover, we conducted a comprehensive comparison with SOTA models, noting that these models were not specifically designed for colonoscopy images and may generate data that hinder the training process or lack the ability for effective improvement. Nevertheless, our experimental results confirm the superiority of our proposed method.

We conducted an ablation study to assess the importance of each proposed component. Table 3.3 and 3.3 report the overall accuracies on the test set. The results demonstrate both components contribute to the accuracy improvement of baseline models, indicating their essential roles in achieving the best final performance.

To further investigate the generative performance of our approach, we present visualization results in Fig. 3, which displays the generated samples and their corresponding masks, alongside the original images for reference. The generated samples demonstrate differences from the original images in both the polyp regions and the backgrounds while maintaining alignment with the masks. Additionally, we sought evaluations from medical professionals to assess the authenticity of the generated samples, and non-medical professionals to locate polyps in the images, which yielded positive feedback on the quality of the generated samples.