DOCSEARCHbyAlgolia
Submit Docs

EfficientSAM Documentation

repository·main·Indexed 25 days ago

https://github.com/yformer/efficientsam

EfficientSAM is a foundation model for efficient segment anything capabilities, utilizing masked image pretraining for high-performance instance segmentation. It supports point-prompt, box-prompt, saliency, and 'segment everything' modes. The library provides two model variants, EfficientSAM-S and EfficientSAM-Ti, with corresponding build functions build_efficient_sam_vits and build_efficient_sam_vitt.

Tokens
1.4K
Snippets
3
Records
7
Agent score
32%

What's inside EfficientSAM

  1. Overview of EfficientSAM capabilities

    main

    EfficientSAM is a model designed for efficient segment anything tasks. It supports several prompting and segmentation modes:

    • Point-prompt: Segmenting objects based on specific point coordinates.
    • Box-prompt: Segmenting objects within a defined bounding box.
    • Segment everything: Automatic segmentation of all objects in an image.
    • Saliency: Identifying and segmenting the most salient parts of an image.

  3. Install and setup EfficientSAM

    main

    To use EfficientSAM, clone the repository and navigate into the directory. You will need to extract the model weights from the provided zip file located in the weights/ directory.

    Required dependencies include torch, torchvision, PIL (Pillow), numpy, and matplotlib.

    !git clone https://github.com/yformer/EfficientSAM.git
os.chdir("EfficientSAM")
    import zipfile
# Extracting weights for EfficientSAM-S
with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
    zip_ref.extractall("weights")

  4. Instantiate EfficientSAM models

    main
    You can instantiate EfficientSAM models using the provided build functions. The repository supports two main model variants: EfficientSAM-S and EfficientSAM-Ti. Use the corresponding build function to load the model with its checkpoints.

  5. Perform segmentation using box or point prompts

    main

    Use the run_ours_box_or_points function to generate segmentation masks. This function accepts an image path, sampled points, point labels, and the model instance.

    Parameters:

    • img_path: Path to the input image.
    • pts_sampled: Coordinates for the prompts (e.g., points or box corners).
    • pts_labels: Labels corresponding to the points (e.g., 1 for positive, 0 for negative).
    • model: The initialized EfficientSAM model.

    Returns:

    • A boolean numpy array representing the predicted mask.
    def run_ours_box_or_points(img_path, pts_sampled, pts_labels, model):
    image_np = np.array(Image.open(img_path))
    img_tensor = ToTensor()(image_np)
    pts_sampled = torch.reshape(torch.tensor(pts_sampled), [1, 1, -1, 2])
    pts_labels = torch.reshape(torch.tensor(pts_labels), [1, 1, -1])
    predicted_logits, predicted_iou = model(
        img_tensor[None, ...],
        pts_sampled,
        pts_labels,
    )

    sorted_ids = torch.argsort(predicted_iou, dim=-1, descending=True)
    predicted_iou = torch.take_along_dim(predicted_iou, sorted_ids, dim=2)
    predicted_logits = torch.take_along_dim(
        predicted_logits, sorted_ids[..., None, None], dim=2
    )

    return torch.ge(predicted_logits[0, 0, 0, :, :], 0).cpu().detach().numpy()

  6. Initialize EfficientSAM models

    main

    EfficientSAM provides different model variants via builder functions. You can initialize the VIT-tiny (build_efficient_sam_vitt) or VIT-small (build_efficient_sam_vits) models. After building, ensure you call .eval() on the model to set it to evaluation mode.

    from efficient_sam.build_efficient_sam import build_efficient_sam_vitt, build_efficient_sam_vits
from squeeze_sam.build_squeeze_sam import build_squeeze_sam

# Initialize VIT-tiny
efficient_sam_vitt_model = build_efficient_sam_vitt()
efficient_sam_vitt_model.eval()

# Initialize VIT-small (requires weights extraction first)
efficient_sam_vits_model = build_efficient_sam_vits()
efficient_sam_vits_model.eval()

# Initialize SqueezeSAM
squeeze_sam_model = build_squeeze_sam()
squeeze_sam_model.eval()

  7. Visualize segmentation masks and prompts

    main

    The following utility functions are used to visualize the results of the segmentation process using matplotlib:

    • show_mask(mask, ax, random_color=False): Overlays the predicted mask on the provided axes.
    • show_points(coords, labels, ax, marker_size=375): Draws prompt points (green stars for positive, red stars for negative).
    • show_box(box, ax): Draws a bounding box prompt.
    • show_anns_ours(mask, ax): A specific visualization style for displaying annotations.