Skip to content

Latest commit

 

History

History
43 lines (31 loc) · 3.8 KB

training.md

File metadata and controls

43 lines (31 loc) · 3.8 KB
Raw

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.