ML: SAM changed the background to white or other colors

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

# https://github.com/facebookresearch/segment-anything/issues/221 import cv2 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/girl.png' img_arr=cv2.imread(inpath) img_arr=cv2.cvtColor(img_arr,cv2.COLOR_BGR2RGB) mask_generator=SamAutomaticMaskGenerator(sam) 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] # plt.imshow(img_arr) # plt.axis('off') # plt.show() filename='D:/PyTest/kkk/girl_new.png' img_arr=cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB) cv2.imwrite(filename, img_arr)