ML: A simple SAM program from input image to output mask, saving as .tif format

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

# https://github.com/facebookresearch/segment-anything/issues/221 import cv2, os import matplotlib.pyplot as plt sam_checkpoint='D:/PyTest/kkk/sam_vit_l_0b3195.pth' model_type="vit_l" from segment_anything import SamAutomaticMaskGenerator, sam_model_registry sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) inpath='D:/PyTest/kkk/dog.jpg' img_arr=cv2.imread(inpath) img_arr=cv2.cvtColor(img_arr,cv2.COLOR_BGR2RGB) mask_generator=SamAutomaticMaskGenerator(sam) # mask_generator = 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, # Requires open-cv to run post-processing # ) predictor=mask_generator.generate(img_arr) # Choose the first mask # mask=predictor[0]['segmentation'] # # Remove background by turn it to white # img_arr[mask==False]=[255, 255, 255] newimg = img_arr[:, :, 0] * 0 for ii in range(len(predictor)): # print(ii) mask=predictor[ii]['segmentation'] # newimg = img_arr[:, :, 0] * 0 newimg[mask == True] = ii+1 # filename = os.path.join('D:/PyTest/kkk/export',str(ii+1)+'.tif') # cv2.imwrite(filename, newimg) # plt.imshow(img_arr) # plt.axis('off') # plt.show() filename='D:/PyTest/kkk/dog_new.tif' cv2.imwrite(filename, newimg)