EfficientSam封装API

CVPR 2024 满分论文！Meta提出EfficientSAM：快速分割一切！
EfficientSAM GitHub

EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything
EfficientSAM采用蒸馏的方法，将模型压缩，虽然精度上比对SAM有所欠缺，但是进行加速了图像编码的速度。官方虽然进行了示例代码，但是还是感觉不好用，所以又进一步封装了代码。主要包含prompt 模式和Everthing模式

Everthing

from eitsam_process.efficientsam_api import (
    get_efficient_sam_model,
    EfficientSAMPrompt,
    EfficientSAMEverthing
) 

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = get_efficient_sam_model(gpu=DEVICE)

box_generate = EfficientSAMPrompt(gpu,model)
input_point = np.array([[500,200], [750, 550]]) #[[x1, y1],[x2, y2]]

# input_label is used to indicate whether it is a point or a bounding box
# [[0,0]] is a positive point, [[1.1]] is a negative point, [[2,3]] is box
input_label = np.array([[2,3]])
masks = box_generate.segment_prompt(input_point, input_label, image_path="img.jpg")  

Prompt

import cv2
import numpy as np

everthing_generate = EfficientSAMEverthing(grid_size=16,gpu=DEVICE, model=model)
masks = everthing_generate.segment_everthing(image_path="img.jpg")
print(len(masks))
for i,mask in enumerate(masks):
    if mask["area"] < 1000: continue
    cv2.imwrite(f"imgs/sub_img{i}.png", np.uint8(mask["segmentation"]*255))

• 通过上面几行代码即可选择不同的提示方法进行分割得到掩码。
• 完整的代码可在https://github.com/FangWenRE/EfficientSam-API.git查看。