SAMCT

This repo is the official implementation for:
SAMCT: Segment Any CT Allowing Labor-Free Task-Indicator Prompts.
(The details of our SAMCT can be found in the models directory in this repo or the paper.)

Highlights

🏆 SAMCT supports two modes: interactive segmentation and automatic segmentation. (allowing both manual prompts and labor-free task-indicator prompts)
🏆 CT5M is a large CT dataset. (about 1.1M images and 5M masks covering 118 categories)
🏆 Excellent performance. (superior to both foundation models and task-specific models)

Installation

Following Segment Anything and SAMUS, python=3.8.16, pytorch=1.8.0, and torchvision=0.9.0 are used in SAMCT.

1. Clone the repository.

   git clone https://github.com/xianlin7/SAMCT.git
   cd SAMCT
   

2. Create a virtual environment for SAMCT and activate the environment.

   conda create -n SAMCT python=3.8
   conda activate SAMCT
   

3. Install Pytorch [pytorch=1.8.0] and TorchVision [torchvision=0.9.0]. (you can follow the instructions here)
4. Install other dependencies.

  pip install -r requirements.txt

(* If you have already installed our SAMUS, you can skip steps 2-4 above, and activate the environment of SAMUS conda activate SAMUS)

Checkpoints

• We use checkpoint of SAM in vit_b version during training SAMCT.
• The checkpoint of SAMCT trained on CT5M will be released in the future 🌝.

Data

• CT5M consists of 30 publicly-available datasets:
  • Hemorrhage-BrainCTImages
  • INSTANCE
  • HaN-Seg
  • Totalsegmentator
  • HaN-OAR
  • PDDCA
  • Covid-19-20
  • MosMed
  • MM-WHS
  • DecathlonZ Task06
  • DecathlonZ Task07
  • DecathlonZ Task09
  • DecathlonZ Task10
  • COVID-19 CT Scan
  • FUMPE
  • LCTSC
  • COVID-19 CT Image Seg
  • LNDB
  • ATM
  • Pediatric-CT-SEG
  • DecathlonZ Task08
  • AMOS22
  • WORD
  • BTCV
  • LiTS
  • CHAOS
  • KiTS23
  • NIH-Pancreas
  • Adrenal-ACC-Ki67
  • Prostate-Anatomical-Edge-Cases
• All images were saved in PNG format. No special pre-processed methods are used in data preparation.
• We have provided some examples to help you organize your data. Please refer to the file fold example_of_required_dataset_format.
  Specifically, each line in all_train/all_val.txt should be formatted as follows:

    <class ID>/<dataset file folder name>/<image file name>
  

  (Here, "class ID" represents the label value of each category in the indicated dataset. For example, the "class ID" of the spleen, right kidney, left kidney, gallbladder, and esophagus on the BTCV dataset should be 1, 2, 3, 4, and 5, respectively.)
• The relevant information of your data should be configured in ./utils/config.py.
• We will release the processed CT5M in the future 🌝.

Training

Once you have the data ready, you can start training the model.

cd "/home/...  .../SAMCT/"
python train.py --modelname SAMCT --task <your dataset config name>
python train_auto_prompt.py --modelname AutoSAMCT --task <your dataset config name>

Testing

Do not forget to set the load_path in ./utils/config.py before testing.

python testSAMCT.py --modelname SAMCT --task <your dataset config name>
python test.py --modelname AutoSAMCT --task <your dataset config name>

Citation

If our SAMCT is helpful to you, please consider citing:

@misc{lin2024samct,
      title={SAMCT: Segment Any CT Allowing Labor-Free Task-Indicator Prompts}, 
      author={Xian Lin and Yangyang Xiang and Zhehao Wang and Kwang-Ting Cheng and Zengqiang Yan and Li Yu},
      year={2024},
      eprint={2403.13258},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}