training.md

Guideline for training SegVol

Build universal datasets

📂 If the 25 processed datasets in our ModelScope/魔搭社区 or HuggingFace have been downloaded, you can skip this step.

1. We use the Abdomenct-12organ as demo dataset.
2. After downloading the demo dataset, you need to config the script/build_dataset.sh file to set the environment vars:
   • $SAVE_ROOT is the save path for the post-processed datasets.
   • $DATASET_CODE is your custom id for your dataset. We suggest you use 0000, 0001, ... as the dataset id.
   • $IMAGE_DIR and $LABEL_DIR is the image directory path and label directory path of the original demo dataset.
   • $TEST_RATIO is the ratio of preserved val/test data from the whole set.
3. Set the category in data_process/train_data_process.py file. Categories should be in the same order as the corresponding idx in ground truth volume and background category should be ignored.
4. Just run bash script/build_dataset.sh.

If you want to combine multiple datasets, you can run the script/build_dataset.sh for multiple times and assign different $DATASET_CODE for each dataset.

Build pseudo mask labels

After the process of building universal datasets finished, you should build pseudo mask labels for each CT in the post-processed datasets.

1. You will need to config the script/build_pseudo_mask.sh first:
   • $DATASET_ROOT is the directory path for the post-processed datasets.
   • $DATASET_CODE is the custom code of your post-processed dataset.
2. Run bash script/build_pseudo_mask.sh. The pseudo masks for the $DATASET_CODE dataset will be generated at $DATASET_ROOT/$DATASET_CODE/fh_seg.

If you combine multiple datasets, you should run the script/build_pseudo_mask.sh for each dataset.

Training

1. Make sure you have completed the above steps correctly.
2. Set environment vars in script/train.sh:
   • $SEGVOL_CKPT is the weight file of SegVol.(Download from huggingface/BAAI/SegVol🤗 or Google Drive)
   • $WORK_DIR is save path for log files and checkpoint files in the training phase.
   • $DATA_DIR is the directory path for the above post-processed datasets.
   • Define dataset_codes to indicate which datasets are used for training
   • Configure these parameters according to your training needs.
   • Set the $CUDA_VISIBLE_DEVICES according to your devices.
3. Run bash script/train.sh.

Training from scratch

If you want to training from scratch without our SegVol checkpoint, I highly recommend that you use the pre-trained ViT here and modify here to load the CLIP TextEncoder parameters.