Learning 3D Ultrasound Segmentation under Extreme Label Deficiency

Environment Setup

Install pytorch following the official guide. Then install other dependencies:

pip install numpy opencv-python tqdm

Download the pretrained DINO model and the sample CETUS dataset from this folder and unzip. The total setup time should be less than 1h.

Train the teacher network using only 0.1% of the ground truth labels. The training takes less than 12h in a single V100 GPU.

python train_teacher.py --workers 32 --batchsize 32 --anno_ratio 0.001 --epochs 100

Use the teacher network to generate high-quality 3D labels:

python generate_pseudo_anno.py --sequence SEQ_PATH --method randaugopt --head VIT_HEAD_PATH

Train the student network on the generated labels. The training takes less than 24h in a workstation with 4 V100 GPUs.

python train_student.py --workers 32 --batchsize 16 --data data/cetus --anno_ratio 0.001 --epochs 100 --pseudo_method randaugopt

Use the student network to perform 3D segmentation:

python inference.py --method student_kd --pretrained STUDENT_CKPT --input_dir IMG_DIR --output_dir OUTPUT_DIR