mazurowski-lab / SegmentAnyBone

SegmentAnyBone: A Universal Model that Segments Any Bone at Any Location on MRI
https://arxiv.org/abs/2401.12974
Other
66 stars 15 forks source link
bone-segmentation deep-learning medical-imaging segment-anything-model segmentation

SegmentAnyBone: A Universal Model that Segments Any Bone at Any Location on MRI

arXiv Paper

SegmentAnyBone is a foundational model based bone segmentation algorithm adapted from Segment Anything Model (SAM) for MRI scans. It is able to segment bones in the following 17 body parts:

Warning: Please note that this software is developed for research purposes and is not intended for clinical use yet. Users should exercise caution and are advised against employing it immediately in clinical or medical settings.

Humerus | Thoracic Spine | Lumbar Spine | Forearm | Pelvis | Hand | Lower Leg

Shoulder | Chest | Arm | Elbow | Hip | Wrist | Thigh | Knee | Foot | Ankle

Screenshot

Dataset

Screenshot

Installation & Usage

You can clone the repository and install required Python packages by running following commands:

git clone https://github.com/mazurowski-lab/SegmentAnyBone.git
cd SegmentAnybone; pip install -r requirements.txt

Sample Output

Screenshot

Model Checkpoints

You can download required model checkpoints from following links: Mobile SAM SegmentAnyBone

After cloning the repository and downloading the checkpoints to the project folder, you should put your 3D MRI volume in /images, and your ground truth mask in /masks folder if you want to evaluate the segmentation performance of SegmentAnyBone. If you need to segment 3D volume instead of 1 slice you can see the 3D segmentation mask under /predicted_masks after you run predictVolume() or predictAndEvaluateVolume() function. This notebook will guide you to use SegmentAnyBone in a slice-based and volume-based manner thanks to following predictSlice(), evaluateSlicePrediction()_, and predictAndEvaluateVolume() functions:

ori_img, predictedSliceMask, atten_map = predictSlice(
    image_name = '2.nii.gz', 
    lower_percentile = 1,
    upper_percentile = 99,
    slice_id = 50, # slice number
    attention_enabled = True, # if you want to use the depth attention
)

msk_gt, dsc_gt = evaluateSlicePrediction(
    mask_pred = predictedSliceMask, 
    mask_name = '2.nrrd', 
    slice_id = 50
)
mask = predictVolume(
    image_name = '2.nii.gz', 
    lower_percentile = 1, 
    upper_percentile = 99
)

predictAndEvaluateVolume(
    image_name = '2.nii.gz', 
    mask_name = '2.nrrd',
    lower_percentile = 1, 
    upper_percentile = 99
)

License

The model is licensed under the CC BY-NC 4.0

Citation

If you find our work to be useful for your research, please cite our paper:

@misc{gu2024segmentanybone,
      title={SegmentAnyBone: A Universal Model that Segments Any Bone at Any Location on MRI}, 
      author={Hanxue Gu and Roy Colglazier and Haoyu Dong and Jikai Zhang and Yaqian Chen and Zafer Yildiz and Yuwen Chen and Lin Li and Jichen Yang and Jay Willhite and Alex M. Meyer and Brian Guo and Yashvi Atul Shah and Emily Luo and Shipra Rajput and Sally Kuehn and Clark Bulleit and Kevin A. Wu and Jisoo Lee and Brandon Ramirez and Darui Lu and Jay M. Levin and Maciej A. Mazurowski},
      year={2024},
      eprint={2401.12974},
      archivePrefix={arXiv},
      primaryClass={eess.IV}

}