Is there a function to generate masks fully automatically, like the notebook in the ex

Hi <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="

Hi Thanks for the scrip <a class="user-mention notranslate" data-hovercard-type="u

Thanks! We also provide a <a href="https://colab.research.google.com/drive/1dhRq4eR6Fb

"Automatic mask generation" is possible? about sam-hq HOT 5 OPEN

syscv commented on September 1, 2024

"Automatic mask generation" is possible?

from sam-hq.

Comments (5)

benji2264 commented on September 1, 2024 1

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