Open mattyamonaca opened 1 year ago
Hi @mattyamonaca, you can already do it with exactly the same code as SAM :)
from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
# Read image
image = cv2.imread("dog.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Load HQ-SAM
sam_checkpoint = "sam_hq_vit_h.pth"
model_type = "vit_h"
device = "cuda"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
# Generate segmentation
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(image)
This gives you a list of masks. Each mask is a dictionary for which the "segmentation"
key gives you a boolean segmentation mask. You can visualize them with the show_anns()
function provided by SAM:
def show_anns(anns):
if len(anns) == 0:
return
sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
ax = plt.gca()
ax.set_autoscale_on(False)
img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
img[:,:,3] = 0
for ann in sorted_anns:
m = ann['segmentation']
color_mask = np.concatenate([np.random.random(3), [0.35]])
img[m] = color_mask
ax.imshow(img)
plt.figure(figsize=(12,7))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()
Thanks!
Hi Thanks for the scrip @benji2264 . When i run the code in jupyter, i get the following error:
TypeError: 'SamAutomaticMaskGenerator' object is not callable
any insight?
cheers!
Hi @SuroshAhmadZobair, thank you for pointing this out, it should work now :) I had forgot to call the generate()
function
Is there a function to generate masks fully automatically, like the notebook in the example? Also, do you have any plans to add such a function?
https://github.com/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb