# segment anything from efficientvit.sam_model_zoo import create_efficientvit_sam_model from efficientvit.models.efficientvit.sam import EfficientViTSamAutomaticMaskGenerator import cv2 import matplotlib.pyplot as plt import numpy as np import os def write_masks_to_folder(masks, path: str) -> None: header = "id,area,bbox_x0,bbox_y0,bbox_w,bbox_h,point_input_x,point_input_y,predicted_iou,stability_score,crop_box_x0,crop_box_y0,crop_box_w,crop_box_h" # noqa metadata = [header] for i, mask_data in enumerate(masks): mask = mask_data["segmentation"] filename = f"{i}.png" cv2.imwrite(os.path.join(path, filename), mask * 255) mask_metadata = [ str(i), str(mask_data["area"]), *[str(x) for x in mask_data["bbox"]], *[str(x) for x in mask_data["point_coords"][0]], str(mask_data["predicted_iou"]), str(mask_data["stability_score"]), *[str(x) for x in mask_data["crop_box"]], ] row = ",".join(mask_metadata) metadata.append(row) metadata_path = os.path.join(path, "metadata.csv") with open(metadata_path, "w") as f: f.write("\n".join(metadata)) return def show_anns(anns): if len(anns) == 0: return sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True) 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.85]]) img[m] = color_mask plt.imshow(img) path = "./1.jpg" image = cv2.imread(path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) plt.figure(figsize=(20,20)) plt.imshow(image) plt.axis('off') efficientvit_sam = create_efficientvit_sam_model(name="efficientvit-sam-xl1", pretrained=True) efficientvit_sam = efficientvit_sam.cpu().eval() efficientvit_mask_generator = EfficientViTSamAutomaticMaskGenerator(efficientvit_sam) masks = efficientvit_mask_generator.generate(image) write_masks_to_folder(masks, "./output")# 在背景图的基础上直接覆盖分割图 show_anns(masks) plt.show()