使用 ultralytics库的SAM / FastSAM 进行图像分割的步骤

1. 环境准备

首先，确保你已经安装了所需的库：

pip install opencv-python ultralytics

2. 加载模型

加载 SAM 模型的代码如下：

from ultralytics import FastSAM
from ultralytics import SAM

model = FastSAM("FastSAM-s.pt") #使用fastsam
#model = SAM("FastSAM-s.pt")  #使用sam

3. 图像分割

通过模型对输入图像进行分割，并保存结果：

results = model("/media/lindsay/data/ultralytics-main/sam.jpg", save=True)
for r in results:
    print(f"Detected {len(r.masks)} masks")

4. 提取掩码数据

接下来，我们提取掩码数据，并将其转换为 NumPy 数组，便于后续处理：

masks = results[0].masks.data.cpu().numpy()

5. 保存和显示掩码

我们可以迭代每个掩码，保存为图像文件，并选择性地在窗口中显示：

import cv2

for index in range(len(masks)):
    mask_image = (masks[index] * 255).astype('uint8')
    cv2.imwrite(f"mask_image_{index}.png", mask_image)

6. 完整代码

将以上步骤整合在一起，完整代码如下：

import cv2
from ultralytics import FastSAM

# Load model
model = FastSAM("FastSAM-s.pt")

# Segment with point prompt
results = model("/media/lindsay/data/ultralytics-main/sam.jpg", save=True)
for r in results:
    print(f"Detected {len(r.masks)} masks")

# Extract and save masks
masks = results[0].masks.data.cpu().numpy()
for index in range(len(masks)):
    mask_image = (masks[index] * 255).astype('uint8')
    cv2.imwrite(f"mask_image_{index}.png", mask_image)