Towards Universal Text-driven CT Image Segmentation

Deep learning-based methods [1, 2] have been widely applied for organ segmentation and tumor segmentation and detection, yielding promising results. However, these methods are often task-specific or organ-specific, such as organ segmentation [25] or tumor detection [26, 27]. Recently, there has been a growing interest in building foundation segmentation models for various organs and tumors [13, 19, 28], and for this reason adapting SAM for ’universal’ segmentation [15, 29]; however, SAM’s interactive nature still requires considerable manual input, limiting its applicability in clinical settings [30]. Additionally, since SAM was developed using natural images, it may lack the medical semantic understanding to differentiate between healthy organs and tumors. In contrast, our method integrates medical professional’s knowledge from radiology reports into a text model, which is more practical than visual prompting for clinical usage.

Vision-language models such as CLIP [20] have demonstrated the ability to learn transferable visual features through language supervision, without manual image annotations. In medical imaging, where diagnostic radiology-founded reports complement imaging data, CLIP-based models have shown promise across various tasks, including organ segmentation [19], disease classification [22], and image-text retrieval [31]. However, adapting CLIP to medical imaging presents unique challenges. Unlike natural images, which can often be summarized in a few sentences, medical images like CT scans contain complex diagnostic information, resulting in detailed radiology reports. This complexity demands improved local-level alignment between image and text. Another significant challenge is data scarcity. While CLIP is ideal for large-scale datasets, medical image-text datasets are relatively limited, which requires a more label-efficient approach. Our method addresses this by curating a radiology-specific corpus of image-text pairs, designed to supplement the radiology-specific knowledge that conventional CLIP training lacks.

Prompt engineering has demonstrated strong performance improvements in both natural language processing and computer vision tasks. By utilizing pre-trained vision-language models [20], text prompting methods results in open-vocabulary segmentation [32, 33] and referring segmentation [34] tasks. Previous studies demonstrate that pretrained CLIP models are effective for 2D medical image segmentation tasks [35, 17]. Li et al. propose a text-augmented segmentation model that utilizes medical text annotations and pseudo labels in a semi-supervised framework [24]. However, extending these approaches to 3D vision-language models for CT segmentation remains an active research area due to the limited availability of paired CT image-text datasets. Recent efforts have centered on developing text-prompted universal models for segmenting various organs and tumors in 3D volumes. Liu et al. proposed to leverage CLIP’s text embeddings to guide the segmentation model for partially-labeled datasets [19]. Zhao et al. introduced SAT, a large-scale segmentation model with a knowledge-enhanced text encoder for multimodal organ and tumor segmentation [16]. However, no existing approach fully leverages a vision-language-aligned model like CLIP for 3D medical image segmentation tasks. Given CLIP’s strong grounding capabilities, we argue that vision-language alignment is essential for building robust text-driven segmentation models. To address this, we propose a pre-trained vision-language model trained on a large-scale CT image-report dataset, incorporating diverse captions for each organ and region.

In recent years, language supervision methods such as CLIP [20] and SigLIP [36] have been shown to be effective in learning transferable representations. Concretely, given a batch of images {$I_{1}$,$I_{2}$,···,$I_{N}$} and paired text descriptions {$T_{1}$,$T_{2}$,···,$T_{N}$}, CLIP leverages an image encoder $E_{i}$ to extract the image embedding $v_{N}$ and a text encoder $E_{t}$ to extract the text embedding $t_{N}$:

The extracted image embedding and text embedding are used to compute InfoNCE loss, in which paired image-text samples are considered as positives and the unpaired ones as negatives. The image-to-text loss ${L}_{i2t}$ can be formulated as:

where cos⟨·,·⟩ denotes the cosine similarity and ${\tau}$ is a learnable temperature parameter. The final bidirectional total loss can be formulated as $\mathcal{L}_{\text{CLIP}}=\frac{1}{2}(\mathcal{L}^{i2t}_{\text{CLIP}}+\mathcal{L}^{t2i}_{\text{CLIP}})$.

While CLIP has demonstrated promising results on natural images using large-scale datasets, it still faces considerable challenges in medical data due to the significant domain gap between them. First, paired medical image and text data are much more limited compared to natural images, necessitating the use of image-level and text-level data augmentations. Second, medical text typically comes in the form of diagnostic reports made by radiologists, which are much longer and more complex than the concise captions found in natural image datasets. While diagnostic reports contain abundant information about the patient, they may lack the granular, organ-level details essential for effective text-driven segmentation. Also, directly training the model with these lengthy reports may cause it to overgeneralize, attempting to encompass all organs and abnormalities in a single image without distinguishing among them.

Recognizing these challenges, we propose to leverage existing LLMs to generate granular text descriptions relating to organs and abnormalities from each CT image’s radiology report. As in-context learning shows great promise in aligning LLM’s responses to human gold standards, we also leverage this technique to generate organ-level descriptions. First, we generate some few-shot examples of concise organ-level captions from the radiology impressions (Figure 1). We then prompt GPT-4 and Llama-3 APIs with ”You are a medical expert. You will read radiology reports from various physicians on CT images. Given the following radiology report, generate shorter captions describing each organ or disease.” to generate target texts. The LLMs will receive the few-shot examples as context and the radiology report as input data, returning the generated organ-level descriptions. By using in-context learning, we ensure the model maintains consistency in its response and preserves the semantic details in the original caption. Finally, observing that the generated captions may contain redundant information, we followed MetaCLIP’s [37] approach to filter out low-quality captions via substring matching. We utilize a radiology-specific text corpus RadLex, which provides an agreed set of named entities for radiology procedures, as our text metadata. We then apply sub-string matching on the granular captions with the metadata entries, which identifies high-quality captions that contain any of the metadata entries, filtering the various types of noise that the LLM may introduce. By addressing these challenges, we aim to enhance the training process and expand the benefits of augmentation to the text inputs, leading to improved performance and a more comprehensive learning framework.

Having generated and filtered the image captions, we can improve the image-text alignment via a multi-granularity contrastive learning framework. The main difference between our approach and CLIP’s is that our approach leverages multiple granular captions as text augmentation. Given the granular captions $C^{l}=[c_{1},\dots,c_{M}]$, we implement a simple random sampling strategy to gather short captions:

where $S_{i,j}$ refers to the $j$-th caption of $i$-th sample from the short caption set. The multi-granular contrastive text-to-image loss becomes:

where $\mathbf{t}_{i,j}$ refers to the text embedding of $S_{i,j}$, and $K$ denotes the number of randomly sampled short captions (K=3 in our implementation). We formulate the bi-directional multi-granular contrastive loss: $\mathcal{L}_{\text{MGCL}}=\left(\mathcal{L}^{t2v}_{\text{MGCL}}+\mathcal{L}^{v2t}_{\text{MGCL}}\right)/2$. By introducing short captions as text augmentations, we increase the diversity of the training data, resulting in a more complete and aligned representation of both images and text.

Empirically, we find that image-text pretraining may not result in optimal finetuning performance for dense prediction tasks like segmentation. To enhance the model’s dense prediction capability, we first pretrained the image encoder on TotalSegmentator for segmentation for 500 epochs, providing a good initialization for the image encoder. We then initialized the image encoder with the pretrained model’s weights before image-text alignment. During image-text contrastive learning, we locked the image encoder and only tuned the text encoder to align the text embeddings with the vision embeddings.

For the final pretraining objective, we combine the original CLIP loss $\mathcal{L}_{\text{CLIP}}$ using the full radiology report with the proposed $\mathcal{L}_{\text{MGCL}}$ loss using the generated short captions:

This approach helps the model to capture detailed connections between images and text while keeping the broader context provided in the radiology reports, improving its ability to generalize to a variety of image-text pairs.

We compare our method with both vision-only segmentation models and text-driven segmentation frameworks. For vision-only methods, we compare with nnUNetv2 [2] and UniMiSS [28]. For text-driven approaches, we compare with CLIP-Driven Universal Model [19], SAT-Pro [16], and CT-CLIP [10]. As shown in Table 1, our method is able to achieve better performance than both the previous SOTA image-only methods and the text-driven SOTA methods on TotalSegmentator dataset. Compared to the best image-only method UniMISS, our method outperforms by 2.7% DSC. Compared to the best text-driven method, i.e., SAT, our method outperforms by 3.1% DSC. We also visualize the organ segmentation results in (Figure 3). This shows that our method can effectively segment the majority of organs with superior performance. For tumor segmentation tasks, as shown in Table 2, our method achieves comparable performance as the best text-driven method Universal model and outperforms the best vision-only method nnUNet by 4.3% on average.

Compared to vision-only models, text-driven segmentation models are more flexible by parsing a wide range of clinical descriptions to guide the segmentation process. This allows text-driven models to generalize to partially labeled data that are incomplete or inconsistent as compared to training data. In real-world scenarios where healthcare professionals use varied prompts to describe or annotate specific organs, this generalization capability becomes particularly valuable for achieving universal organ segmentation. To assess these generalization capabilities, we evaluate how well the text-driven segmentation model handles two categories of training invisible prompts: 1) prompts obtained by merging multiple organs and 2) prompts that are synonymous terms for the target organ. Specifically, for category 1, we obtain these prompts by merging various suborgans used in training (i.e. left lung is obtained by merging the lung upper lobe left, lung lower lobe left classes in TotalSegmentator). For category 2, we take the training visible prompt and substitute it for a synonym. For example, renal organs is synonym for kidney and hepatic system is synonym for liver.

Generalization results to merged suborgans is shown in Table 3. Compared to the CLIP-Driven Universal Model, SAT Pro, and CT-CLIP, our method consistently achieves superior performance on merging simple left and right organs (e.g. left and right lungs, left and right kidneys). These results highlight the model’s ability to interpret novel combinations of suborgans effectively. Generalization results to synonyms is shown in Table 4. Notably, our method also achieves significantly higher performance in challenging cases (e.g., cervical vertebrae, lumbar vertebrae, thoracic vertebrae, veins, arteries) without explicitly being trained on such prompts. Our method consistently outperforms the existing text-driven methods, achieving the highest average performance (73.2% DSC) across all categories.

We visualize the segmentation results of generalization study in axial view (Figure 5) and coronal view (Figure 6). As shown in Figure 5 rows 1, 3, and 4, our model performs resonably well when merging left and right organs in both the chest region (lung) and the abdominal region (kidney, autochthon). In row 2, the model demonstrates its flexibility by accurately segmenting organs described using synonyms, such as ”hepatic system” for the liver and ”renal organs” for the kidneys. Additionally, for bones and vertebrae, our method effectively segments merged categories like left ribs, right ribs, lumbar vertebrae, and thoracic vertebrae (Figure 5). This demonstrates the superior ability of our model to handle diverse and unseen clinical terminology as text prompts, suitable for real-world deployment in diverse clinical environments.

We further study the generalization capabilities to external datasets. Table 6 summarizes the performance on FLARE22 dataset. Compared to the other methods, our method consistently achieves superior generalization performances in 10 of 13 organs in the abdominal region. The second best performing model is the vision-only model nnUNet, which our method slightly outperforms by 0.4% DSC. Generalization performance for SegTHOR dataset is shown in table 7. Our method also achieves superior performance in esophagus, heart and aorta segmentation. For inferring the training unseen heart category, we prompt text-driven models with the prompt heart. Our method outperforms best existing text-driven models in heart segmentation by 9.4% DSC and 6% DSC on average. To further explore the generalization capability of nnUNet on the unseen heart category, we combined predictions for its sub-organ components (i.e., heart myocardium, heart ventricle, heart atrium, and pulmonary artery). Interestingly, nnUNet demonstrates comparable performance on this cardaic organ category, achieving results superior to most existing text-driven models (except ours). We hypothesize that this is because text-driven models may suffer from insufficient image-text alignment during training, limiting their ability to generalize effectively to unseen categories (heart in this case).