DiffULD: Diffusive Universal Lesion Detection

Universal Lesion Detection can be formulated as locating lesions in input CT scan $I_{ct}$ with a set of boxes predictions ${z^{{}^{\prime}}_{0}}$. For a particular box ${z^{{}^{\prime}}}$, it can be denoted as ${z^{{}^{\prime}}}=[{x}_{1},{y}_{1},{x}_{2},{y}_{2}]$, where ${x}_{1},{y}_{1}$ and ${x}_{2},{y}_{2}$ are the coordinates of the top-left and bottom-right corners, respectively.

We design our model based on a diffusion model mentioned in [40]. As shown in Fig. 2, our method consists of two stages, a forward diffusion (or training) process and a reverse refinement (or inference) process. In the forward process, We denote GT bounding boxes as ${z}_{0}$ and generate corrupted training samples ${z}_{1},{z}_{2},...,{z}_{T}$ for latter DiffULD training by adding Gaussian noise iteratively, which can be defined as:

where $\bar{\alpha}_{t}$ represents the noise variance schedule and $t\in\{0,1,...,T\}$. Subsequently, a neural network ${f}_{\theta}(z_{t},t,I_{ct})$ conditioned on the corresponding CT scan $I_{ct}$ is trained to predict $z_{0}$ from a noisy box ${z}_{T}$ by reversing the noising process step by step. During inference, for an input CT scan $I_{ct}$ with a set of random boxes, the model is able to refine the random boxes to get a lesion detection prediction box ${z^{{}^{\prime}}_{0}}$, iteratively.

In this section, we specify the training process with our novelty introduced ‘Center-aligned BBox padding’.

Center-aligned BBox padding. As shown in Fig. 1, we utilize Center-aligned BBox padding to generate perturbated boxes. Then, the perturbated boxes are paired with original GT BBoxes, forming perturbated GT boxes which are used to generate corrupted training samples ${z}_{1},{z}_{2},...,{z}_{T}$ for the latter DiffULD training by adding Gaussian noise iteratively.

For clarity, we denote $[x_{c}^{i},y_{c}^{i}]$ as center coordinates of original GT BBox, where $z^{i}$, $[w^{i},h^{i}]$ are the width and height of $z^{i}$.

We consider the generation in two parts: box scaling and center sampling. (i) Box scaling: We set a hyper-parameter $\lambda_{scale}\in(0,1)$ for scaling. For $z^{i}$, The width and height of the corresponding perturbated boxes are randomly sampled in uniform distributions on $U_{w}\sim[(1-\lambda_{scale})w^{i},(1+\lambda_{scale})w^{i}]$ and $U_{h}\sim[(1-\lambda_{scale})h^{i},(1+\lambda_{scale})h^{i}]$. (ii) Center sampling: We sample the center coordinates $[{x}_{c}^{new},{y}_{c}^{new}]$ of perturbated boxes from a 2D Gaussian distribution $\mathcal{N}$ whose probability density function can be denoted as:

where $\sigma$ is a size-adaptive parameter, which can be calculated according to the $z^{i}$’s width and height:

Besides, for each input CT scan $I_{ct}$, we collect all GT BBoxes in ${z}$ and add random perturbations to them and generate multiple perturbated boxes for each of them. Thus the number of perturbated boxes in an image varies with its number of GT BBoxes. For better training, we fix the number of perturbated boxes as $N$ for all training images.

As shown in Fig 1., the perturbated boxes cluster together and their centers are still aligned with the corresponding original GT BBox. Subsequently, perturbated GT boxes $z_{0}$ are sent for corruption as the training objective.

Box corruption. As shown in Fig. 1, we corrupt the parameters of $z_{0}$ with Gaussian noises. The noise scale is controlled by $\bar{\alpha}_{t}$ (in Eq. 1), which adopts the decreasing cosine scheduler in the different time step $t$.

Loss function. As the model generates the same number of ($N$) predictions for the input image, termed as a prediction set, the loss function should be set-wised [34]. Specifically, each GT is matched with the prediction by the least matching cost, and the overall training loss[34] can be represented as:

where $\mathcal{L}_{L1box}$ and $\mathcal{L}_{giou}$ are the pairwise box loss. We adopt $\lambda_{L1box}=2.0$ and $\lambda_{giou}=5.0$.

At the inference stage, with a set of random boxes sampled from Gaussian distribution, the model does refinement step by step to obtain the final predictions $z^{{}^{\prime}}_{0}$. For better performance, two key components are used:

Box filtering. In each refinement step, the model receives a set of box proposals from the last step. As the prediction starts from arbitrary boxes and the lack of GT (lesion), most of these proposals are very far from lesions. Keeping refining them in the following steps will hinder network training. Toward efficient detection, we send the proposals to the detection head and remove the boxes whose confidential scores are lower than a particular threshold $\lambda_{conf}$. The remaining high-quality proposals are sent for followed DDIM sampling.

Box update with DDIM sampling. DDIM[41] is utilized to further refine the received box proposals by denoising. Next, these refined boxes are sent to the next step and start a new round of refinement. After multiple steps, final predictions are obtained.

However, we observe that if we just filter out boxes with low scores during iterative refinement, the model runs out of usable box proposals rapidly, which also leads to a deterioration in performance. Therefore, after the box updating, we add new boxes sampled from a Gaussian Distribution to the set of remaining boxes with. The number of box proposals per image is padded to the fixed number $N$ before they are sent to the next refinement step.

Multi-window input. Most prior arts in ULD use a single and fixed window (e.g., a wide window of [1024, 4096]) to render the input CT scan, which suppresses organ-specific information and makes it hard for the network to focus on the various lesions. Therefore, taking cues from [8], we introduce 3 organ-specific HU windows to highlight multiple organs of interest. Their window widths and window levels are: ${W}_{1}=[1200,-600]$ for chest organs, ${W}_{2}=[400,50]$ for soft tissues and ${W}_{3}=[200,30]$ for abdomenal organs. Multi-window features are extracted with a ConvNeXt-T shared network.

3D context feature fusion. We modify the original A3D[16] DenseNet backbone for context fusion. We remove the first Conv3D Block and use the truncated network as our 3D context fusion module, which fuses the multi-view features from the last module. Multi-window features are fused with this module and sent to the detector subsequently for lesion detection.

Our experiments are conducted on the standard ULD dataset DeepLesion[1]. The dataset contains 32,735 lesions on 32,120 axial slices from 10,594 CT studies of 4,427 unique patients. We use the official data split of DeepLesion which consists of 70%, 15%, 15% for training, validation, and test, respectively. Besides, we also evaluate the performance of 3 methods based on a revised test set from .

Training details. DiffULD is trained on CT scans of size 512 × 512 with a batch size of 4 on 4 NVIDIA RTX Titan GPUs with 24GB memory. For hyper-parameters, the threshold $N$ for box padding is set to $300$. $\lambda_{scale}$ for box scaling is set to 0.4. $\lambda_{conf}$ for box filtering is set to 0.5. We use the Adam optimizer with an initial learning rate of $2e-4$ and the weight decay as $1e-4$. The default training schedule is 120K iterations, with the learning rate divided by 10 at 60K and 100K iterations. Data augmentation strategies contain random horizontal flipping, rotation, and random brightness adjustment.

Evaluation metrics. The lesion detection is classified as true positive (TP) when the IoU between the predicted and the GT bounding box is larger than 0.5. Average sensitivities computed at 0.5, 1, 2, and 4 false-positives (FP) per image are reported as the evaluation metrics on the test set for a fair comparison.

We evaluate the effectiveness of DiffULD against anchor-based ULD approaches such as 3DCE[7], MVP-Net[8], A3D[16] and SATr[31] on DeepLesion. Several anchor-free natural image detection methods such as FCOS[33] and DN-DETR[44] are also introduced for comparison. In addition, we conduct an extensive experiment to explore DiffULD’s potential on an augmented test set of completely annotated DeepLesion volumes, introduced by Lesion-Harvester[45].

Table 1. demonstrates that our proposed DiffULD achieves favorable performances when compared to recent state-of-the-art anchor-based ULD approaches such as SATr on both 3 slices and 7 slices. It outperforms prior well-established methods such as A3D and MULAN by a non-trivial margin. This validates that with our padding strategy, the concise DPM can be utilized in general medical object detection tasks such as ULD and attain impressive performance.

We provide an ablation study about our proposed approach: center-aligned BBox padding. As shown in Table 3., we compared it with various other padding strategies, including: (i) duplicating original GT boxes; (ii) padding random boxes sampled from a uniform distribution; (iii) padding random boxes sampled from a Gaussian distribution; (iv) padding with center-aligned strategy.

Our baseline method is training the diffusion model directly with no box padding, using our proposed backbone in 2.4. The performance is increased by 0.30% by simply duplicating the original GT boxes. Padding random boxes following uniform and Gaussian distributions brings 0.51% and 0.91% improvement respectively. Our center-aligned padding strategy accounts for 1.13% improvement from the baseline. We attribute this performance boost to center-aligned padding’s ability to provide high-quality additional training targets. It effectively expands the insufficient training targets on CT scans and enhances DPM’s detection performance while avoiding deterioration triggered by adding random targets. This property is favorable for utilizing DPMs on a limited amount of GT like ULD.