🚀 SlimSAM(SAM壓縮版,即“任意分割模型”)
SlimSAM是Segment Anything(SAM)模型的壓縮(剪枝)版本,能夠根據點或框等輸入提示生成高質量的對象掩碼。它通過創新的壓縮方法,在大幅降低訓練成本的同時,實現了接近原始模型的性能。
🚀 快速開始
提示掩碼生成
from PIL import Image
import requests
from transformers import SamModel, SamProcessor
model = SamModel.from_pretrained("nielsr/slimsam-50-uniform")
processor = SamProcessor.from_pretrained("nielsr/slimsam-50-uniform")
img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
input_points = [[[450, 600]]]
inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to("cuda")
outputs = model(**inputs)
masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
scores = outputs.iou_scores
在生成掩碼時,除了其他參數,你可以傳入感興趣對象的大致二維位置、包圍感興趣對象的邊界框(格式應為邊界框右上角和左下角的x、y座標)、分割掩碼。截至撰寫本文時,根據官方倉庫,官方模型不支持將文本作為輸入。更多詳細信息,請參考這個筆記本,它通過可視化示例展示瞭如何使用該模型!
自動掩碼生成
該模型可用於以“零樣本”方式根據輸入圖像生成分割掩碼。模型會自動用一個包含1024
個點的網格進行提示,並將這些點全部輸入模型。
以下是自動掩碼生成的示例代碼:
from transformers import pipeline
generator = pipeline(task="mask-generation", model="nielsr/slimsam-50-uniform", device = 0, points_per_batch = 256)
image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
outputs = generator(image_url, points_per_batch = 256)
以下是顯示圖像的代碼:
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
def show_mask(mask, ax, random_color=False):
if random_color:
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
else:
color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
ax.imshow(mask_image)
plt.imshow(np.array(raw_image))
ax = plt.gca()
for mask in outputs["masks"]:
show_mask(mask, ax=ax, random_color=True)
plt.axis("off")
plt.show()
✨ 主要特性
SlimSAM是一種新穎的SAM壓縮方法,通過統一的剪枝-蒸餾框架高效重用預訓練的SAM,實現了卓越的性能和極低的訓練成本。具體特性如下:
- 高效壓縮:通過創新的交替瘦身策略,將壓縮過程劃分為漸進式步驟,在大幅減少參數和計算量的同時,保持接近原始模型的性能。
- 低訓練成本:與其他現有方法相比,訓練成本降低了10倍以上,僅需0.1%(10k)的SAM訓練數據。
- 高性能表現:在參數數量減少至僅0.9%(570萬)、MACs減少至0.8%(21G)的情況下,仍能實現接近原始SAM-H的性能。
📚 詳細文檔
模型詳情
SAM模型由3個模塊組成:
VisionEncoder
:基於VIT的圖像編碼器。它使用注意力機制對圖像塊進行計算,以生成圖像嵌入,並使用相對位置嵌入。
PromptEncoder
:為點和邊界框生成嵌入。
MaskDecoder
:一種雙向變壓器,在圖像嵌入和點嵌入之間進行交叉注意力計算,並將輸出結果輸入到Neck
模塊。
Neck
:根據MaskDecoder
生成的上下文掩碼預測輸出掩碼。
📄 許可證
本項目採用Apache-2.0許可證。
📜 引用
如果您使用此模型,請使用以下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}
}
@misc{chen202301,
title={0.1% Data Makes Segment Anything Slim},
author={Zigeng Chen and Gongfan Fang and Xinyin Ma and Xinchao Wang},
year={2023},
eprint={2312.05284},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
模型相關圖片
SlimSAM概述及其與其他替代方案的差異。
原倉庫鏈接
點擊查看原倉庫
免責聲明
本模型卡片的內容由Hugging Face團隊撰寫,部分內容從原始的SAM模型卡片複製粘貼而來。