Segment Anything 工具

Segment Anything 工具使用 Segment Anything Model (SAM)。 此文件夹包含代表 SAM 的 ONNX 文件，供浏览器用于计算分割。

获取模型检查点

SAM 由 Facebook 提供，形式为 PyTorch 模型：一个 .pth 文件。 模型有三种类型。从大到小依次为 vit_h、vit_l 和 vit_b。 建议使用 vit_b，因为它更小、更快，代价是分割质量稍低。

转换为 ONNX

Facebook 提供了将解码器转换为 ONNX 的脚本，但没有提供编码器的转换脚本。维护者拒绝合并添加此功能的 PR。 你可以使用这些 PR 中提供的脚本，但我们选择的解决方案是使用 samexporter。

• 创建一个文件夹，用于存放转换脚本和上一步下载的原始 checkpoint.pth 模型。

mkdir temp
cd temp

• 克隆 samexporter 仓库和 segment anything 仓库

git clone git@github.com:vietanhdev/samexporter.git
git clone git@github.com:facebookresearch/segment-anything.git

• 使用 pip 安装 segment anything 以及 samexporter 的依赖（在 python 3.11.5 的虚拟环境中测试通过）

pip install -e ./segment-anything
pip install torchvision==0.16.1 onnx==1.15.0 onnxruntime==1.15.1 timm==0.9.12

• 进入 samexporter 文件夹并运行命令导出编码器和解码器（不要使用量化）

cd samexporter
python -m samexporter.export_encoder --checkpoint ../checkpoint.pth --output ../encoder.onnx --model-type vit_b
python -m samexporter.export_decoder --checkpoint ../checkpoint.pth --output ../decoder.onnx --model-type vit_b --return-single-mask

• 将编码器和解码器复制到项目中的正确位置（目前只有 vit_b 模型可用）

cd $NIMUS_IMAGE_DIR
cp encoder.onnx public/onnx-models/sam/$MODEL_NAME
cp decoder.onnx public/onnx-models/sam/$MODEL_NAME