多伦多大学Wangbo博士实验室 LiteMedSAM

https://github.com/bowang-lab/MedSAM/blob/LiteMedSAM/README.md

恭喜Techna科学家Wangbo博士，他最近被任命为加拿大CIFAR AI主席(link is external)的第二届AI主席(link is external)中的一员。

Wangbo博士领导Peter Munk心脏中心（PMCC）的人工智能团队，致力于开发机器学习算法，帮助预测疾病的发展趋势，并协助临床医生做出个性化患者护理的决策。这一卓越的CIFAR主席将支持并推动这项工作的发展。

除了在Techna和PMCC的现任职务外，Wangbo博士还被交叉任命为多伦多大学医学生物物理学系和Vector研究所的成员。他毕业于斯坦福大学计算机科学系，并曾在多家生物技术公司，包括Illumina和Genentech，担任人工智能科学家。

CIFAR AI主席支持加拿大在人工智能研究和培训方面的领导地位，为每位主席提供高达100万加元的资金支持。它们是由CIFAR管理的全加拿大人工智能战略(link is external)的一部分，该战略是对全国人工智能研究社区进行1.25亿加元投资的一部分。

---

MedSAM的轻量级版本，用于快速训练和推理。该模型经过以下两个阶段的训练：

阶段1. 通过将图像嵌入输出设为相同，从MedSAM图像编码器ViT中提炼出轻量级图像编码器TinyViT。

阶段2. 用TinyViT替换MedSAM图像编码器ViT，并对整个流程进行微调。

安装步骤：

1. 使用以下环境进行测试：Ubuntu 20.04 | Python 3.10 | CUDA 11.8 | PyTorch 2.1.2

2. 创建虚拟环境：

   conda create -n medsam python=3.10 -y
   

3. 激活虚拟环境：

   conda activate medsam
   

4. 安装 PyTorch 2.0。

5. 克隆代码库：

   git clone -b LiteMedSAM https://github.com/bowang-lab/MedSAM/
   

6. 进入MedSAM文件夹：

   cd MedSAM
   

7. 运行以下命令安装依赖：

   pip install -e .
   

快速教程：提交到CVPR 2024 MedSAM on Laptop挑战

健全性测试

1. 下载LiteMedSAM检查点并将其放置在work_dir/LiteMedSAM目录中。
2. 下载演示数据。
3. 执行以下命令进行健全性测试：

python CVPR24_LiteMedSAM_infer.py -i test_demo/imgs/ -o test_demo/segs

构建Docker

docker build -f Dockerfile -t litemedsam .

注意：别忘了在最后加上 .

在测试演示图像上运行Docker

docker container run -m 8G --name litemedsam --rm -v $PWD/test_demo/imgs/:/workspace/inputs/ -v $PWD/test_demo/litemedsam-seg/:/workspace/outputs/ litemedsam:latest /bin/bash -c "sh predict.sh"

注意：如果遇到“Permission denied”错误，请运行 chmod -R 777 ./*。

保存Docker

docker save litemedsam | gzip -c > litemedsam.tar.gz

计算指标

python evaluation/compute_metrics.py -s test_demo/litemedsam-seg -g test_demo/gts -csv_dir ./metrics.csv

模型训练

数据预处理

1. 下载 Lite-MedSAM 检查点并将其放置在当前目录下。
2. 下载演示数据集。本教程假设已将其解压缩到 data/FLARE22Train/ 目录下。
3. 运行预处理脚本将数据集转换为 npz 格式：

python pre_CT_MR.py \
    -img_path data/FLARE22Train/images \ ## 训练图像路径
    -img_name_suffix _0000.nii.gz \ ## 训练图像扩展名
    -gt_path data/FLARE22Train/labels \ ## 训练标签路径
    -gt_name_suffix .nii.gz \ ## 训练标签扩展名
    -output_path data \ ## 保存预处理数据的路径
    -num_workers 4 \ ## 预处理的工作线程数
    -prefix CT_Abd_ \ ## 预处理数据的前缀
    -modality CT \ ## 预处理数据的模态
    -anatomy Abd \ ## 预处理数据的解剖结构
    -window_level 40 \ ## CT 的窗位
    -window_width 400 \ ## CT 的窗宽
    --save_nii ## 同时以nii.gz格式保存预处理数据，以便在其他软件中进行可视化检查

数据集划分

将演示数据集的前40个案例用于训练，保存在 MedSAM_train 中，将最后10个用于测试，保存在 MedSAM_test 中。

有关脚本的详细使用方法，请运行 python pre_CT_MR.py -h。

将训练的npz文件转换为npy格式

python npz_to_npy.py \
    -npz_dir data/MedSAM_train \ ## 预处理npz训练数据的路径
    -npy_dir data/npy \ ## 保存用于训练的转换后的npy数据的路径
    -num_workers 4 ## 并行转换的工作线程数

微调预训练的 Lite-MedSAM

训练流程需要约10GB的GPU内存，批量大小设置为4。

单GPU

要在单个GPU上训练 Lite-MedSAM，运行：

python train_one_gpu.py \
    -data_root data/MedSAM_train \
    -pretrained_checkpoint lite_medsam.pth \
    -work_dir work_dir \
    -num_workers 4 \
    -batch_size 4 \
    -num_epochs 10

要从检查点恢复中断的训练，请运行：

python train_one_gpu.py \
    -data_root data/MedSAM_train \
    -resume work_dir/medsam_lite_latest.pth \
    -work_dir work_dir \
    -num_workers 4 \
    -batch_size 4 \
    -num_epochs 10

有关其他命令行参数，请运行 python train_one_gpu.py -h。

多GPU

要在多个GPU上微调 Lite-MedSAM，运行：

python train_multi_gpus.py \
    -i data/npy \ ## 训练数据集的路径
    -task_name MedSAM-Lite-Box \
    -pretrained_checkpoint lite_medsam.pth \
    -work_dir ./work_dir_ddp \
    -batch_size 16 \
    -num_workers 8 \
    -lr 0.0005 \
    --data_aug \ ## 使用数据增强
    -world_size <WORLD_SIZE> \ ## 将要使用的GPU总数
    -node_rank 0 \ ## 如果在单台机器上训练，则设置为0
    -init_method tcp://<MASTER_ADDR>:<MASTER_PORT>

或者，您可以使用提供的 train_multi_gpus.sh 脚本在多个GPU上进行训练。要从检查点恢复中断的训练，请在多GPU训练的命令行参数中添加 -resume <your_work_dir>，而不是检查点路径；脚本将自动在工作目录中找到最新的检查点。有关其他命令行参数，请运行 python train_multi_gpus.py -h。

推断（健全性测试）

推断脚本假设测试数据已转换为npz格式。要在3D CT FLARE22数据集上运行推断，运行：

python inference_3D.py \
    -data_root data/npz/MedSAM_test/CT_Abd \ ## 预处理的npz数据
    -pred_save_dir ./preds/CT_Abd \
    -medsam_lite_checkpoint_path work_dir/medsam_lite_latest.pth \
    -num_workers 4 \
    --save_overlay \ ## 保存分割叠加在输入图像上
    -png_save_dir ./preds/CT_Abd_overlay \ ## 仅在设置了 --save_overlay 时使用
    --overwrite ## 覆盖现有预测，默认从现有预测继续

有关其他命令行参数，请运行 python inference_3D.py -h。

我们还提供一个在2D图像上运行推断的脚本 inference_2D.py，其用法与3D脚本相同。

常见问题解答（FAQ）

1. 预处理的npz数据和npy数据之间有什么区别？

   npz格式用于存储2D和3D图像（关注ROI）以及它们对应的地面实况掩模。在3D图像的情况下，间距也被包括在内。所有这些数据元素都被紧凑地打包到一个npz文件中。这种格式主要用于分发我们的训练和验证数据集。值得注意的是，对于验证数据集，提供了边界框而不是地面实况标签掩模。

   相比之下，npy格式将每个2D图像或3D图像的单个切片以及其标签掩模分别存储在不同的文件中。在加载数据进行训练时使用此格式。

2. 我在加载已训练模型的检查点进行推断时遇到问题，应该怎么办？

   如果在加载已训练模型的检查点进行推断时遇到困难，我们建议用户首先尝试使用位于 MedSAM/utils/ 下的 extract_weights.py 脚本。此脚本用于从现有检查点中提取权重并将其保存到新的检查点文件中。

   要使用此脚本，请在终端中执行以下命令：

   python extract_weights.py \
       -from_pth <YOUR_CHECKPOINT_PATH> \
       -to_pth <NEW_CHECKPOINT_PATH>
   

   将 <YOUR_CHECKPOINT_PATH> 替换为保存的已训练模型检查点的路径，将 <NEW_CHECKPOINT_PATH> 替换为新检查点文件的所需路径。执行此命令并创建新的检查点后，应该可以在推断任务中使用。

致谢

我们感谢MobileSAM和TinyViT的作者公开提供其源代码。