Segment Anything

Meta AI Research, FAIR

Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alex Berg, Wan-Yen Lo, Piotr Dollar, Ross Girshick

[论文] [项目] [演示] [数据集] [博客] [BibTeX]

Segment Anything模型 (SAM) 从输入提示（如点或框）生成高质量的对象掩码，并可用于生成图像中所有对象的掩码。它已在包含1100万张图像和11亿个掩码的数据集上进行了训练，并在各种分割任务上具有强大的零样本性能。

免责声明

我不是原始实现的作者。这个存储库是一个Python包，它将原始实现包装成一个包。如果您对原始实现有任何疑问，请参考facebookresearch/segment-anything。原始实现采用Apache License，版本2.0。

安装

代码需要python>=3.8，以及pytorch>=1.7和torchvision>=0.8。请按照此处的说明安装PyTorch和TorchVision依赖项。强烈建议安装支持CUDA的PyTorch和TorchVision。

使用pip安装Segment Anything

pip install segment-anything-py

或使用mamba安装Segment Anything

mamba install -c conda-forge segment-anything

或克隆本地存储库并使用以下命令安装

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

以下可选依赖项对于掩码后处理、将掩码保存为COCO格式、示例笔记本以及将模型导出为ONNX格式是必要的。jupyter也是运行示例笔记本所必需的。

pip install opencv-python pycocotools matplotlib onnxruntime onnx

入门

首先下载一个模型检查点。然后，只需几行代码就可以使用模型从给定的提示中获取掩码

from segment_anything import SamPredictor, sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)

或为整张图片生成掩码

from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(<your_image>)

此外，还可以从命令行生成图像的掩码

python scripts/amg.py --checkpoint <path/to/checkpoint> --model-type <model_type> --input <image_or_folder> --output <path/to/output>

有关更多详细信息，请参阅使用提示与SAM一起使用和自动生成掩码的示例笔记本。

ONNX导出

SAM的轻量级掩码解码器可以导出为ONNX格式，以便可以在任何支持ONNX运行时环境（如演示中展示的浏览器中）运行。使用以下命令导出模型

python scripts/export_onnx_model.py --checkpoint <path/to/checkpoint> --model-type <model_type> --output <path/to/output>

请参阅示例笔记本，了解如何将图像预处理与SAM的骨干网络结合使用，以及如何使用ONNX模型进行掩码预测。建议使用最新稳定的PyTorch版本进行ONNX导出。

Web演示

demo/文件夹中有一个简单的单页React应用程序，它展示了如何在Web浏览器中使用多线程运行导出的ONNX模型的掩码预测。请参阅demo/README.md以获取更多详细信息。

模型检查点

该模型有三个版本，具有不同的骨干网络大小。可以通过运行以下命令实例化这些模型

from segment_anything import sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")

点击以下链接下载对应模型类型的检查点

• default或vit_h：ViT-H SAM模型。
• vit_l：ViT-L SAM模型。
• vit_b：ViT-B SAM模型。

数据集

请参阅此处了解数据集概述。数据集可以从此处下载。通过下载数据集，您同意已阅读并接受SA-1B数据集研究许可的条款。

我们将每张图片的掩码保存为json文件。在Python中可以按以下格式将其加载为字典。

{
    "image"                 : image_info,
    "annotations"           : [annotation],
}

image_info {
    "image_id"              : int,              # Image id
    "width"                 : int,              # Image width
    "height"                : int,              # Image height
    "file_name"             : str,              # Image filename
}

annotation {
    "id"                    : int,              # Annotation id
    "segmentation"          : dict,             # Mask saved in COCO RLE format.
    "bbox"                  : [x, y, w, h],     # The box around the mask, in XYWH format
    "area"                  : int,              # The area in pixels of the mask
    "predicted_iou"         : float,            # The model's own prediction of the mask's quality
    "stability_score"       : float,            # A measure of the mask's quality
    "crop_box"              : [x, y, w, h],     # The crop of the image used to generate the mask, in XYWH format
    "point_coords"          : [[x, y]],         # The point coordinates input to the model to generate the mask
}

图像ID可以在sa_images_ids.txt中找到，该文件可以使用上述链接下载。

将COCO RLE格式的掩码解码为二进制

from pycocotools import mask as mask_utils
mask = mask_utils.decode(annotation["segmentation"])

请参阅此处获取更多有关操作存储在RLE格式中的掩码的说明。

许可证

该模型采用Apache 2.0许可。

贡献

请参阅贡献和行为准则。

贡献者

“Segment Anything”项目得以实现，离不开众多贡献者（按字母顺序）的帮助。

Adam Adcock, Vaibhav Aggarwal, Morteza Behrooz, Cheng-Yang Fu, Ashley Gabriel, Ahuva Goldstand, Allen Goodman, Sumanth Gurram, Jiabo Hu, Somya Jain, Devansh Kukreja, Robert Kuo, Joshua Lane, Yanghao Li, Lilian Luong, Jitendra Malik, Mallika Malhotra, William Ngan, Omkar Parkhi, Nikhil Raina, Dirk Rowe, Neil Sejoor, Vanessa Stark, Bala Varadarajan, Bram Wasti, Zachary Winstrom

引用“Segment Anything”

如果您在研究中使用了SAM或SA-1B，请使用以下BibTeX条目。

@article{kirillov2023segany,
  title={Segment Anything},
  author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
  journal={arXiv:2304.02643},
  year={2023}
}