如何使用分段任意模型（SAM）保存分段图像或掩模？

Question

我按照教程编写了一个小项目：

import torch 
import cv2 
import matplotlib.pyplot as plt 
import numpy as np 
from pycocotools import mask as mask_util 
import os 
import json 

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) 

def show_points(coords, labels, ax, marker_size=375): 
   pos_points = coords[labels==1] 
   neg_points = coords[labels==0] 
   ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) 
   ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax): 
   x0, y0 = box[0], box[1] 
   w, h = box[2] - box[0], box[3] - box[1] 
   ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2)) 

image = cv2.imread('/home/luisgpm/all_images/1_bcs_1.0.jpeg') 
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(10,10)) 
plt.imshow(image) 
plt.axis('on') 
plt.show()

import sys 
sys.path.append("/home/luisgpm/myenv/lib/python3.11/site-packages/segment_anything/") 
from segment_anything import sam_model_registry, SamPredictor 

sam_checkpoint = "/home/luisgpm/sam_vit_h_4b8939.pth" 
model_type = "default" 
device = "cpu" 

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

predictor = SamPredictor(sam) 
predictor.set_image(image)

input_point = np.array([[500, 375]]) 
input_label = np.array([1]) 
plt.figure(figsize=(10,10)) 
plt.imshow(image) 
show_points(input_point, input_label, plt.gca()) 
plt.axis('on') 
plt.show()  



masks, scores, logits = predictor.predict( 
    point_coords=input_point, 
    point_labels=input_label, 
    multimask_output=True, 
)

masks.shape

for i, (mask, score) in enumerate(zip(masks, scores)): 
    plt.figure(figsize=(10,10)) 
    plt.imshow(image) 
    show_mask(mask, plt.gca()) 
    show_points(input_point, input_label, plt.gca()) 
    plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18) 
    plt.axis('off') 
    plt.show() 

input_box = np.array([80, 35, 795, 1200]) 
masks, _, _ = predictor.predict( 
    point_coords=None, 
    point_labels=None, 
    box=input_box[None, :], 
    multimask_output=False, 
)

plt.figure(figsize=(10, 10)) 
plt.imshow(image) 
show_mask(masks[0], plt.gca()) 
show_box(input_box, plt.gca()) 
plt.axis('off') 
plt.show()

现在我想保存这个分割的结果，它可以是图像本身或掩模，我尝试在github上关注这些问题： https://github.com/facebookresearch/segment-anything/issues/442 https://github.com/facebookresearch/segment-anything/issues/221

但是效果不太好，主要是出现了IndexError，我不知道如何解决，如何以简单的方式保存结果？

Answer 1

您可以使用 OpenCV 和 NumPy 保存分割掩模和掩模图像。具体方法如下：

1. 保存分割蒙版： 您可以将蒙版保存为图像，方法是将其转换为适当的格式，然后使用 cv2.imwrite。

   mask_image = (mask * 255).astype(np.uint8) # 转换为uint8格式 cv2.imwrite('mask.png', mask_image)

2. 保存蒙版图像： 如果你想将蒙版覆盖在原始图像上并保存，可以使用以下代码：

   color_mask = np.zeros_like(图像) color_mask[mask > 0.5] = [30, 144, 255] # 选择你喜欢的颜色 masked_image = cv2.addWeighted(图像, 0.6, color_mask, 0.4, 0)

   cv2.imwrite('masked_image.png', cv2.cvtColor(masked_image, cv2.COLOR_RGB2BGR))

这些代码片段会将蒙版和蒙版图像保存在您当前的工作目录中