CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection

TL;DR

The CLIP-Driven Universal Model is proposed, which incorporates text embedding learned from Contrastive Language-Image Pre-training (CLIP) to segmentation models and is computationally more efficient compared with dataset-specific models, generalized better to CT scans from varying sites, and shows stronger transfer learning performance on novel tasks.

Abstract

An increasing number of public datasets have shown a marked impact on automated organ segmentation and tumor detection. However, due to the small size and partially labeled problem of each dataset, as well as a limited investigation of diverse types of tumors, the resulting models are often limited to segmenting specific organs/tumors and ignore the semantics of anatomical structures, nor can they be extended to novel domains. To address these issues, we propose the CLIP-Driven Universal Model, which incorporates text embedding learned from Contrastive Language-Image Pre-training (CLIP) to segmentation models. This CLIP-based label encoding captures anatomical relationships, enabling the model to learn a structured feature embedding and segment 25 organs and 6 types of tumors. The proposed model is developed from an assembly of 14 datasets, using a total of 3,410 CT scans for training and then evaluated on 6,162 external CT scans from 3 additional datasets. We rank first on the Medical Segmentation Decathlon (MSD) public leaderboard and achieve state-of-the-art results on Beyond The Cranial Vault (BTCV). Additionally, the Universal Model is computationally more efficient (6x faster) compared with dataset-specific models, generalized better to CT scans from varying sites, and shows stronger transfer learning performance on novel tasks.

Figures (11)

• Figure 1: Cosine similarity between CLIP embeddings. The CLIP embedding reveals the intrinsic semantics of the anatomical structures by mapping similar concepts close to each other in the embedding space. For example, "Liver" has a large similarity with "Liver Tumor" and "Hepatic Vessel" (the hepatic vessel returns low-oxygen blood from the liver to the heart, which has a high anatomical relationship with the liver).
• Figure 2: Overview. We have developed a Universal Model from an assembly of 14 public datasets of 3,410 CT scans. In total, 25 organs and 6 types of tumors are partially labeled (detailed in Appendix \ref{['tab:public_dataset']}). To deal with partial labels, Universal Model consists of a text branch and a vision branch (§\ref{['sec:universal_model']}). The official test set of MSD and BTCV are used to benchmark the performance of organ segmentation (§\ref{['sec:strong_challenge_ranking']}) and tumor detection (§\ref{['sec:high_specificity']}). 3D-IRCADb, TotalSegmentator and a large-scale private dataset, consisting of 5,038 CT scans with 21 annotated organs, are used for independent, external validation of model generalizability and transferability (§\ref{['sec:properties']}).
• Figure 3: Benchmark on MSD validation dataset. We compare Universal Model with Swin UNETR tang2022self (previously ranked first on the MSD leaderboard) on 5-fold cross-validation of the MSD dataset. Universal Model achieves overall better segmentation performance and offers substantial improvement in the tasks of segmenting liver tumors (+14%), pancreatic tumors (+8%), and colon tumors (+11%).
• Figure 4: Intra-observer variability. We obtain similar performance between pseudo labels generated by the Universal Model (AI) and annotations performed by two human experts (Dr1,2) on 6 organs. Spleen (Spl), liver (Liv), kidneys (Kid), stomach (Sto), gallbladder (Gall), and pancreas (Pan) can be annotated by AI with a similar intra-observer variability to humans. Examples of pseudo labels and human annotations are provided in Appendix \ref{['fig:pseudo_truth_visualization']}.
• Figure 5: Pancreatic tumor detection. Qualitative visualizations of the proposed Universal Model and five competitive baseline methods. We review the detection results of tumors from smaller to larger sizes (Rows 1--3). When it comes to a CT scan without tumor from other hospitals, the Universal Model generalize well in organ segmentation and does not generate many false positives of tumors (Row 4; §\ref{['sec:high_specificity']}). The visualization of tumor detection in other organs (e.g., liver tumors and kidney tumors) can be found in Appendix Figures \ref{['fig:qualitative_visualization_liver_tumor']}--\ref{['fig:qualitative_visualization_kidney_tumor']}.