计算机视觉| 分割大模型Segment Anything(SAM)从0到1使用
安装anaconda
参考博文 最详细的Anaconda新手安装
pytorch 官网查找合适的版本
需要注意pytorch、CUDA、python、Anaconda的四个版本都要在相互支持的范围内,且能够满足Segment Anything的要求:
The code requires python>=3.8, as well as pytorch>=1.7 and torchvision>=0.8. Please follow the instructions here to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.
可以在pytorch官网选择对应的版本,复制run this command对应的命令。
或者查看之前的版本 以前版本如下-whl文件
执行命令前,需要在 Anaconda Prompt 中先激活环境:
conda activate 环境名
然后执行安装 pytorch 命令
选择合适的版本
- 根据你的 Python 版本和 CUDA 版本,选择相应的 .whl 文件。例如,我们使用的是 Python 3.11 和 CUDA 12.6,可以下载 torch-2.6.0+cu126-cp311-cp311-win_amd64.whl(Windows 系统)或 torch-2.6.0+cu126-cp311-cp311-linux_aarch64.whl(Linux 系统)。
- 同样,对于torchvision,根据Python版本和CUDA12.6,我们可以下载torchvision-0.21.0+cu126-cp311-cp311-win_amd64.whl(Windows 系统)或 torchvision-0.21.0+cu126-cp311-cp311-linux_x86_64.whl(Linux 系统)。
- 同样,对于torchaudio,根据Python版本和CUDA12.6,我们可以下载torchaudio-2.6.0+cu126-cp311-cp311-win_amd64.whl(Windows 系统)或 torchaudio-2.6.0+cu126-cp311-cp311-linux_x86_64.whl(Linux 系统)。
下载文件
点击链接下载 .whl 文件到本地。
我的python版本是3.11 ,在4090服务器上下载
所以我选择三个对应版本是
pytorch=2.1.0 torchvision=0.16.0 torchaudio=2.1.0 cp11
安装.whl文件
安装 .whl 文件:
下载完成后,打开命令提示符或终端,导航到 .whl 文件所在的目录,运行以下命令:
pip install <filename>.whl
将 替换为你下载的 .whl 文件的名称。例如:
pip install torch-2.6.0+cu126-cp311-cp311-win_amd64.whl
安装完成之后,检查一下本地是否能够使用PyTorch。我们通过在本地的Python文件中检查是否能够import torch来判断是否能够使用PyTorch
查看安装包
- conda list命令除了得到当前虚拟环境下的所有已安装的包,还会显示关联环境下的已安装的包; pip
- list命令只显示当前虚拟环境下的所有已安装的包;当关联环境下有某个包,再用conda命令在其他虚拟环境下下载该包时,就会直接关联,无需下载;
- pip list命令得到的结果是conda list得到的结果的子集
下载Segment-anything 源码
Segment Anything github网址
选择一个方式进行下载
使用HTTPS
打开控制台,输入
git clone https://github.com/facebookresearch/segment-anything.git
为项目源码安装依赖库
在下载后的源码目录下框下 输入 cmd,打开 cmd 窗口
在项目源码目录下激活虚拟环境
执行python setup.py install
git clone git@github.com:facebookresearch/segment-anything.git
cd segment-anything; pip install -e .
安装必须依赖
pip install opencv-python pycocotools matplotlib onnxruntime onnx
测试代码及其使用案例
下载model checkpoint
三个模型权重比较大,耐心下载
在源码根目录之下创建checkpoint文件夹,将下载好的模型文件放入文件
运行
添加测试文件夹和图片数据
在代码文件下创建一个stamps文件夹
在stamps下创建如下文件夹
在images文件夹下放入要测试的图片
运行代码
运行scripts文件夹下的amg.py文件
运行代码如下
python <当前文件地址>/amg.py --checkpoint <当前文件地址>/checkpoint/sam_vit_b_01ec64.pth --model-type vit_b --input <当前文件地址>/stamps/images --output <当前文件地址>/stamps/masks
将<当前文件地址>里的内容换成自己的绝对地址
运行后的mask文件包下有除了overlay.png文件外的其他文件。
如果想要生成overlay这样的全景文件
可在amg.py文件下添加如下代码
# 手动生成 Overlay 文件
def create_mask_overlay(image_path, masks, output_path='overlay.png', alpha=0.5):"""生成原始图像 + 掩模叠加的可视化结果Args:image_path: 原始图像路径masks: SAM 输出的掩模列表(每个掩模是二进制 0/1 数组)output_path: 叠加图保存路径alpha: 掩模透明度 (0-1)"""# 读取原始图像image = cv2.imread(image_path)image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转为 RGB# 初始化叠加层(全黑)overlay = np.zeros_like(image)# 将所有掩模合并(逻辑或操作)for mask in masks:mask = mask.astype(np.uint8) * 255 # 0/1 → 0/255colored_mask = cv2.applyColorMap(mask, cv2.COLORMAP_JET) # 伪彩色overlay = cv2.bitwise_or(overlay, colored_mask)# 叠加掩模和原始图像result = cv2.addWeighted(image, 1, overlay, alpha, 0)# 保存结果cv2.imwrite(output_path, cv2.cvtColor(result, cv2.COLOR_RGB2BGR))
在amg.py的main函数下添加如下代码
if output_mode == "binary_mask":os.makedirs(save_base, exist_ok=False)write_masks_to_folder(masks, save_base)# ourselfbinary_masks = [m['segmentation'] for m in masks]overlay_path = os.path.join(save_base, f"{base}_overlay.png")create_mask_overlay(t, binary_masks, output_path=overlay_path, alpha=0.5)else:save_file = save_base + ".json"with open(save_file, "w") as f:json.dump(masks, f)
大功告成!!!恭喜