SAM_MaskGenerator_Test/SAM_MaskGenerator_example.py at master · alejandroolivo/SAM_MaskGenerator_Test · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import numpy as np
import time
import matplotlib.pyplot as plt
import cv2
import sys
import io
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

# params
model_size = "large" # small, medium, large
device = "cuda:0" # cuda:0, cpu
image_name = "cartagena.png"

# load image
image = cv2.imread('image_examples/' + image_name)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

sys.path.append("..")

# download model checkpoint (https://github.com/facebookresearch/segment-anything#model-checkpoints)
if(model_size == "small"):
    sam_checkpoint = "model_checkpoint/sam_vit_b_01ec64.pth"
    model_type = "vit_b"

if(model_size == "medium"):
    sam_checkpoint = "model_checkpoint/sam_vit_l_0b3195.pth"
    model_type = "vit_l"

if(model_size == "large"):
    sam_checkpoint = "model_checkpoint/sam_vit_h_4b8939.pth"
    model_type = "vit_h"

# load model and send to device
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

#load mask generator
mask_generator = SamAutomaticMaskGenerator(sam)

# get time
start = time.time()

masks = mask_generator.generate(image)

# get time
end = time.time()
print('Elapsed time = ' + str((end - start)*1000) + ' ms')

#plot image with all masks overlayed
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)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

plt.figure(figsize=(10,10))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()

#plot first mask
# plt.figure(figsize=(10,10))
# plt.imshow(image)
# plt.imshow(masks[0]['segmentation'], alpha=0.5)
# plt.axis('off')
# plt.show()

# plot all masks
# plt.figure(figsize=(10,10))
# plt.imshow(image)
# for mask in masks:
#     plt.imshow(mask['segmentation'], alpha=0.5)
# plt.axis('off')
# plt.show()