How to save segmented images or masks using Segment Anything Model (SAM)?

I followed the tutorials to code a small project:

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()

now i want to save the result of this segmentation, it can be the image itself or the mask, i tried to follow these issues on github: https://github.com/facebookresearch/segment-anything/issues/442 https://github.com/facebookresearch/segment-anything/issues/221

but that didn't work out well, mainly showing IndexError that I didn't know how to solve, how can i save the results in a simple way?

You can save both the segmentation mask and the masked image using OpenCV and NumPy. Here's how you can do it:

1. Saving the Segmentation Mask: You can save the mask as an image by converting it to an appropriate format and then using cv2.imwrite.

   mask_image = (mask * 255).astype(np.uint8)  # Convert to uint8 format
   cv2.imwrite('mask.png', mask_image)
   

2. Saving the Masked Image: If you want to overlay the mask on the original image and save that, you can use the following code:

   color_mask = np.zeros_like(image)
   color_mask[mask > 0.5] = [30, 144, 255] # Choose any color you like
   masked_image = cv2.addWeighted(image, 0.6, color_mask, 0.4, 0)
   
   cv2.imwrite('masked_image.png', cv2.cvtColor(masked_image, cv2.COLOR_RGB2BGR))
   

These code snippets will save the mask and the masked image in your current working directory.