许可协议: MIT
语言:

• 英语
  任务类型: 图像分割
  标签:
• 医学

SegVol是一款通用且交互式的医学体数据图像分割模型。该模型支持通过点提示、框提示和文本提示进行体积分割。通过在9万例未标注的CT扫描数据和6千例标注CT数据上进行训练，这一基础模型能够支持超过200种解剖结构的识别分割。

论文与代码已开源。

关键词: 3D医学SAM、体数据图像分割

快速开始

环境配置

conda create -n segvol_transformers python=3.8  
conda activate segvol_transformers  

需安装pytorch v1.11.0或更高版本。通过以下命令安装关键依赖：

pip install 'monai[all]==0.9.0'  
pip install einops==0.6.1  
pip install transformers==4.18.0  
pip install matplotlib  

测试脚本

from transformers import AutoModel, AutoTokenizer  
import torch  

# 获取设备  
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  

# 加载模型  
# 若无法连接huggingface.co，可下载仓库后设置from_pretrained路径为本地目录，替换"yuxindu/segvol"  
clip_tokenizer = AutoTokenizer.from_pretrained("yuxindu/segvol")  
model = AutoModel.from_pretrained("yuxindu/segvol", trust_remote_code=True, test_mode=True)  
model.model.text_encoder.tokenizer = clip_tokenizer  
model.eval()  
model.to(device)  
print('模型加载完成')  

# 设置案例路径  
# 可从huggingface yuxindu/segvol的files and versions下载此案例  
ct_path = 'path/to/Case_image_00001_0000.nii.gz'  
gt_path = 'path/to/Case_label_00001.nii.gz'  

# 设置类别（对应标注掩码中的唯一值1,2,3,4...）  
categories = ["肝脏", "肾脏", "脾脏", "胰腺"]  

# 生成npy数据格式  
ct_npy, gt_npy = model.processor.preprocess_ct_gt(ct_path, gt_path, category=categories)  
# 若已下载我们预处理的25个数据集，可直接使用处理好的ct_npy, gt_npy文件跳过此步  

# 通过zoom_transform生成缩放视图  
data_item = model.processor.zoom_transform(ct_npy, gt_npy)  

# 手动添加批次维度  
data_item['image'], data_item['label'], data_item['zoom_out_image'], data_item['zoom_out_label'] = \  
data_item['image'].unsqueeze(0).to(device), data_item['label'].unsqueeze(0).to(device), data_item['zoom_out_image'].unsqueeze(0).to(device), data_item['zoom_out_label'].unsqueeze(0).to(device)  

# 以肝脏为例  
cls_idx = 0  

# 文本提示  
text_prompt = [categories[cls_idx]]  

# 点提示  
point_prompt, point_prompt_map = model.processor.point_prompt_b(data_item['zoom_out_label'][0][cls_idx], device=device)   # 输入无批次维度，输出带批次维度  

# 框提示  
bbox_prompt, bbox_prompt_map = model.processor.bbox_prompt_b(data_item['zoom_out_label'][0][cls_idx], device=device)   # 输入无批次维度，输出带批次维度  

print('提示生成完成')  

# SegVol测试前向传播  
# use_zoom: 使用缩放视图  
# point_prompt_group: 使用点提示  
# bbox_prompt_group: 使用框提示  
# text_prompt: 使用文本提示  
logits_mask = model.forward_test(image=data_item['image'],  
      zoomed_image=data_item['zoom_out_image'],  
      # point_prompt_group=[point_prompt, point_prompt_map],  
      bbox_prompt_group=[bbox_prompt, bbox_prompt_map],  
      text_prompt=text_prompt,  
      use_zoom=False  
      )  

# 计算Dice分数  
dice = model.processor.dice_score(logits_mask[0][0], data_item['label'][0][cls_idx])  
print(dice)  

# 将预测结果保存为nii.gz文件  
save_path='./Case_preds_00001.nii.gz'  
model.processor.save_preds(ct_path, save_path, logits_mask[0][0],   
                           start_coord=data_item['foreground_start_coord'],   
                           end_coord=data_item['foreground_end_coord'])  
print('完成')  

训练脚本

from transformers import AutoModel, AutoTokenizer  
import torch  

# 获取设备  
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  

# 加载模型  
# 若无法连接huggingface.co，可下载仓库后设置from_pretrained路径为本地目录，替换"yuxindu/segvol"  
clip_tokenizer = AutoTokenizer.from_pretrained("yuxindu/segvol")  
model = AutoModel.from_pretrained("yuxindu/segvol", trust_remote_code=True, test_mode=False)  
model.model.text_encoder.tokenizer = clip_tokenizer  
model.train()  
model.to(device)  
print('模型加载完成')  

# 设置案例路径  
# 可从huggingface yuxindu/segvol的files and versions下载此案例  
ct_path = 'path/to/Case_image_00001_0000.nii.gz'  
gt_path = 'path/to/Case_label_00001.nii.gz'  

# 设置类别（对应标注掩码中的唯一值1,2,3,4...）  
categories = ["肝脏", "肾脏", "脾脏", "胰腺"]  

# 生成npy数据格式  
ct_npy, gt_npy = model.processor.preprocess_ct_gt(ct_path, gt_path, category=categories)  
# 若已下载我们预处理的25个数据集，可直接使用处理好的ct_npy, gt_npy文件跳过此步  

# 执行训练变换  
data_item = model.processor.train_transform(ct_npy, gt_npy)  

# 训练示例  
# 手动添加批次维度  
image, gt3D = data_item["image"].unsqueeze(0).to(device), data_item["label"].unsqueeze(0).to(device) # 添加批次维度  

loss_step_avg = 0  
for cls_idx in range(len(categories)):  
    # optimizer.zero_grad()  
    organs_cls = categories[cls_idx]  
    labels_cls = gt3D[:, cls_idx]  
    print(image.shape, organs_cls, labels_cls.shape)  
    loss = model.forward_train(image, train_organs=organs_cls, train_labels=labels_cls)  
    loss_step_avg += loss.item()  
    loss.backward()  
    # optimizer.step()  

loss_step_avg /= len(categories)  
print(f'平均损失 {loss_step_avg}')  

# 保存检查点  
model.save_pretrained('./ckpt')