ML: Automatically generating object masks with SAM

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

# Automatically generating object masks with SAM # https://github.com/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb import numpy as np import torch import matplotlib.pyplot as plt import cv2 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) image=cv2.imread('D:/PyTest/kkk/dog.jpg') image=cv2.cvtColor(image,cv2.COLOR_BGR2RGB) # plt.figure(figsize=(20,20)) # plt.imshow(image) # plt.axis('off') # plt.show() # plt.axis('off') import sys sys.path.append("..") from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor sam_checkpoint='D:/PyTest/kkk/sam_vit_h_4b8939.pth' model_type="vit_h" # device = "cuda" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) # sam.to(device=device) mask_generator = SamAutomaticMaskGenerator(sam) masks=mask_generator.generate(image) print(len(masks)) # print(masks[0].keys()) # plt.figure(figsize=(20, 20)) # plt.imshow(image) # show_anns(masks) # plt.axis('off') # plt.show() mask_generator_2 = SamAutomaticMaskGenerator( model=sam, points_per_side=32, pred_iou_thresh=0.86, stability_score_thresh=0.92, crop_n_layers=1, crop_n_points_downscale_factor=2, min_mask_region_area=100, ) mask2=mask_generator_2.generate(image) print(len(mask2)) # plt.figure(figsize=(20, 20)) # plt.imshow(image) # show_anns(mask2) # plt.axis('off') # plt.show()