【计算机视觉】Segment Anything 安装配置及代码测试（含源代码）

一、前言

目前代码已经开源！

Segment Anything Model（SAM）可以从输入提示（如点或框）生成高质量的物体遮罩，并且可以用于为图像中的所有物体生成遮罩。它在一个包含1100万张图像和10亿个遮罩的数据集上进行了训练，并且在各种分割任务上表现出了强大的零样本性能。

二、安装

2.1 基本要求

该代码要求使用 python>=3.8，并且需要安装 pytorch>=1.7 和 torchvision>=0.8。请按照以下说明安装 PyTorch 和 TorchVision 的依赖项。强烈建议同时安装支持 CUDA 的 PyTorch 和 TorchVision。

以下是安装步骤的一般指南：

安装 Python 3.8+：确保您的系统已安装 Python 3.8 或更高版本。您可以从 Python 官方网站（https: //www.python.org/downloads/）下载并安装适用于您的操作系统的 Python 版本。
安装 PyTorch 和 TorchVision：按照以下步骤安装 PyTorch 和 TorchVision：
访问 PyTorch 官方网站（https: //pytorch.org/）并根据您的系统选择适当的安装选项。

根据提供的安装说明，使用 pip 或 conda 安装 PyTorch 和 TorchVision。例如，如果您使用 pip，可以执行以下命令安装 PyTorch：

pip install torch>=1.7 torchvision>=0.8

安装 CUDA（可选）：如果您的系统支持 NVIDIA GPU 并且您希望使用 CUDA 加速，建议安装 CUDA 并配置 PyTorch 和 TorchVision 以支持 CUDA。您可以从 NVIDIA 官方网站（https: //developer.nvidia.com/cuda-downloads）下载适用于您的系统的 CUDA 版本，并按照提供的说明进行安装。

请注意，上述步骤提供了一般的安装指南。具体的安装步骤可能因您的操作系统、Python 版本和其他依赖项而有所不同。建议参考 PyTorch 和 TorchVision 的官方文档和安装说明，以确保正确地安装和配置这些库。

2.2 Install Segment Anything

pip install git+https://github.com/facebookresearch/segment-anything.git

若是这个运行失败，选择下面的方式：

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

cd segment-anything

pip install -e .

便可顺利安装成功！

以下是用于遮罩后处理、以 COCO 格式保存遮罩、示例笔记本和以 ONNX 格式导出模型的可选依赖项。同时，运行示例笔记本还需要安装 jupyter。

pip install opencv-python pycocotools matplotlib onnxruntime onnx

For mask post-processing: You may need to install additional libraries or packages depending on the specific post-processing techniques used in the code. It is recommended to refer to the code documentation or instructions for the required dependencies.

For saving masks in COCO format: If you intend to save the generated masks in COCO format, you will need to install the pycocotools library. You can install it using pip:

pip install pycocotools
pip install jupyter
pip install onnx onnxruntime

三、代码使用示例

3.1 Automatically generating object masks with SAM

from IPython.display import display, HTML
display(HTML(
"""
<a target="_blank" href="https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb">
  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
</a>
"""
))

上述代码片段是用于在Jupyter Notebook或支持HTML输出的环境中显示一个带有Colab徽章的链接。当点击该链接时，它将在Colab中打开名为"automatic_mask_generator_example.ipynb"的笔记本。

要使用此代码片段，请确保已经安装并正确配置了IPython和Jupyter Notebook。将代码片段放置在代码单元格中并运行，您将在输出中看到一个带有Colab徽章的链接，点击该链接即可在Colab中打开相应的笔记本。

3.2 Environment Set-up

using_colab = False

if using_colab:
    import torch
    import torchvision
    print("PyTorch version:", torch.__version__)
    print("Torchvision version:", torchvision.__version__)
    print("CUDA is available:", torch.cuda.is_available())
    import sys
    !{sys.executable} -m pip install opencv-python matplotlib
    !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'
    
    !mkdir images
    !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/dog.jpg
        
    !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

如果在本地使用 Jupyter 运行，请首先根据存储库中的安装说明在您的环境中安装segment_anything。

如果在 Google Colab 上运行，请在下方将 using_colab=True 设置为 True 并运行该单元格。在 Colab 中，请确保在“编辑”->“笔记本设置”->“硬件加速器”下选择了“GPU”。

3.3 显示标注

import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:,:,3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.concatenate([np.random.random(3), [0.35]])
        img[m] = color_mask
    ax.imshow(img)

这段代码是一个用于显示标注（annotations）的函数 show_anns。下面是对代码的解读：

函数接受一个标注列表 anns 作为参数。
首先，检查标注列表的长度，如果列表为空，则直接返回。
根据标注的面积对标注进行排序，从大到小，使用 sorted 函数和 key 参数来实现排序。排序后的结果保存在 sorted_anns 列表中。
创建一个坐标轴对象 ax，并关闭其自动缩放功能。
创建一个图像数组 img，形状与最大标注的分割形状相同，并初始化为全1，表示完全透明。
遍历排序后的标注列表，对每个标注进行处理：
获取标注的分割掩码 m。

生成一个随机的颜色掩码 color_mask，由3个随机数和一个透明度值组成。

将颜色掩码应用到图像数组的相应位置上，使得标注区域显示为对应的颜色。

使用 ax.imshow 函数显示图像数组 img，即显示了带有颜色标注的图像。
总体来说，该函数的作用是根据给定的标注信息，在图像上显示带有不同颜色的标注区域。

3.4 图像示例

image = cv2.imread('images/dog.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(20,20))
plt.imshow(image)
plt.axis('off')
plt.show()

原图如下：

3.5 Automatic mask generation

要运行自动 mask 生成，请向 SamAutomaticMaskGenerator 类提供一个 SAM 模型。将下面的路径设置为 SAM 检查点的路径。推荐在 CUDA 上运行，并使用默认模型。

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"

device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

mask_generator = SamAutomaticMaskGenerator(sam)

masks = mask_generator.generate(image)

print(len(masks))
print(masks[0].keys())

输出结果为：

dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])

plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show() 

可视化结果如下：

3.6 Automatic mask generation options

自动掩码生成中有几个可调参数，用于控制采样点的密度以及去除低质量或重复掩码的阈值。 此外，生成可以在图像的裁剪上自动运行以提高较小对象的性能，并且后处理可以去除杂散像素和孔洞。 以下是对更多掩码进行采样的示例配置：

mask_generator_2 = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,
    pred_iou_thresh=0.86,
    stability_score_thresh=0.92,
    crop_n_layers=1,
    crop_n_points_downscale_factor=2,
    min_mask_region_area=100,  # Requires open-cv to run post-processing
)

masks2 = mask_generator_2.generate(image)

plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks2)
plt.axis('off')
plt.show()