EfficientSAM

https://github.com/yformer/efficientsam
Admin
EfficientSAM is an efficient segment anything model that leverages masked image pretraining to...
Tokens:12,870
Snippets:56
Trust Score:6.6
License:Apache-2.0
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install segment-anything package

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

Installs the segment-anything library from its GitHub repository. This command is typically run in a shell or notebook environment.

```bash
!pip install git+https://github.com/facebookresearch/segment-anything.git
```

--------------------------------

### EfficientSAM Image Segmentation Example

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

This Python code demonstrates how to load an image, run EfficientSAM for segmentation, and display the original image alongside the segmentation results. Ensure necessary libraries like matplotlib, numpy, and PIL are installed, and the EfficientSAM model is loaded.

```python
fig, ax = plt.subplots(1, 2, figsize=(30, 30))
image_path = "figs/examples/dogs.jpg"
image = np.array(Image.open(image_path))
ax[0].imshow(image)
ax[0].title.set_text("Original")
ax[0].axis('off')

ax[1].imshow(image)
mask_efficient_sam_vits = run_everything_ours(image_path, efficient_sam_vits_model)
show_anns_ours(mask_efficient_sam_vits, ax[1])
ax[1].title.set_text("EfficientSAM")
ax[1].axis('off')


plt.show()
```

--------------------------------

### Run EfficientSAM with Point Input (VIT-tiny)

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

This code snippet demonstrates running EfficientSAM with point prompts. It defines input points and labels, loads an image, and then uses `run_ours_box_or_points` to generate segmentation masks. The results for EfficientSAM (VIT-tiny) and EfficientSAM (VIT-small) are displayed.

```python
fig, ax = plt.subplots(1, 4, figsize=(30, 30))

input_label = np.array([1, 1])
image_path = "figs/examples/dogs.jpg"
input_point = np.array([[580, 350], [650, 350]])
image = np.array(Image.open(image_path))
show_points(input_point, input_label, ax[0])
ax[0].imshow(image)


ax[1].imshow(image)
mask_efficient_sam_vitt = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vitt_model)
show_anns_ours(mask_efficient_sam_vitt, ax[1])
ax[1].title.set_text("EfficientSAM (VIT-tiny)")
ax[1].axis('off')

ax[2].imshow(image)
mask_efficient_sam_vits = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vits_model)
show_anns_ours(mask_efficient_sam_vits, ax[2])
```

--------------------------------

### Build and Initialize EfficientSAM Models

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Imports model building functions and initializes EfficientSAM models (vitt and vits variants). Sets the models to evaluation mode.

```python
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
import zipfile

efficient_sam_vitt_model = build_efficient_sam_vitt()
efficient_sam_vitt_model.eval()
```

--------------------------------

### Run EfficientSAM with Box Input (VIT-tiny)

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

This code snippet demonstrates how to run the EfficientSAM model with a bounding box input. It sets up the input points and labels, loads an image, and then calls the `run_ours_box_or_points` function to generate masks. The resulting masks are then displayed.

```python
x1=400
y1=200
x2=800
y2=600
w=x2-x1
h=y2-y1

fig, ax = plt.subplots(1, 4, figsize=(30, 30))
input_point = np.array([[x1, y1], [x2, y2]])
input_label = np.array([2,3])
image_path = "figs/examples/dogs.jpg"
image = np.array(Image.open(image_path))
show_points(input_point, input_label, ax[0])
show_box([x1,y1,x2,y2], ax[0])
ax[0].imshow(image)


ax[1].imshow(image)
mask_efficient_sam_vitt = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vitt_model)
show_anns_ours(mask_efficient_sam_vitt, ax[1])
ax[1].title.set_text("EfficientSAM (VIT-tiny)")
ax[1].axis('off')

ax[2].imshow(image)
mask_efficient_sam_vits = run_ours_box_or_points(image_path, input_point, input_label, efficient_sam_vits_model)
show_anns_ours(mask_efficient_sam_vits, ax[2])
ax[2].title.set_text("EfficientSAM (VIT-small)")
ax[2].axis('off')


ax[3].imshow(image)
mask_squeeze_sam = run_ours_box_or_points(image_path, input_point, input_label, squeeze_sam_model)
show_anns_ours(mask_squeeze_sam, ax[3])
ax[3].title.set_text("SqueezeSAM")
ax[3].axis('off')

plt.show()
```

--------------------------------

### Build and Load EfficientSAM Models

Source: https://context7.com/yformer/efficientsam/llms.txt

Instantiates pre-trained EfficientSAM models (ViT-Tiny for speed, ViT-Small for accuracy). Ensure checkpoints are extracted if zipped. Models are returned in evaluation mode.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt, build_efficient_sam_vits
import zipfile

# Build EfficientSAM-Ti (ViT-Tiny) - faster inference
efficient_sam_ti = build_efficient_sam_vitt()
efficient_sam_ti.eval()

# For EfficientSAM-S, extract the checkpoint from zip first (file is >100MB)
with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
    zip_ref.extractall("weights")

# Build EfficientSAM-S (ViT-Small) - higher accuracy
efficient_sam_s = build_efficient_sam_vits()
efficient_sam_s.eval()

print(f"EfficientSAM-Ti loaded: {sum(p.numel() for p in efficient_sam_ti.parameters())} parameters")
print(f"EfficientSAM-S loaded: {sum(p.numel() for p in efficient_sam_s.parameters())} parameters")
```

--------------------------------

### Clone Repository

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

Downloads the EfficientSAM repository from GitHub.

```bash
!git clone https://github.com/yformer/EfficientSAM.git
import os
os.chdir("EfficientSAM")
```

--------------------------------

### Load EfficientSAM Model Checkpoint

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Extracts the EfficientSAM model checkpoint from a zip file and builds the model instance. Ensure the 'weights' directory exists.

```python
with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
    zip_ref.extractall("weights")
efficient_sam_vits_model = build_efficient_sam_vits()
efficient_sam_vits_model.eval()

squeeze_sam_model = build_squeeze_sam()
squeeze_sam_model.eval()
```

--------------------------------

### Instantiate EfficientSAM Model

Source: https://github.com/yformer/efficientsam/blob/main/README.md

Use these functions to build the EfficientSAM model variants. Ensure the appropriate checkpoints are available.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt, build_efficient_sam_vits
efficientsam = build_efficient_sam_vitt()
```

--------------------------------

### Batch Processing Multiple Images with EfficientSAM

Source: https://context7.com/yformer/efficientsam/llms.txt

Demonstrates loading multiple images, preparing them as a single batch, and running inference with batched point prompts.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt
from PIL import Image
from torchvision import transforms
import torch
import numpy as np
import os

model = build_efficient_sam_vitt()
model.eval()

# Prepare batch of images
image_paths = ["figs/examples/dogs.jpg"] * 4  # Example with same image
batch_size = len(image_paths)

# Load and preprocess images to same size
target_size = (512, 512)
images = []
for path in image_paths:
    img = Image.open(path).resize(target_size)
    img_tensor = transforms.ToTensor()(np.array(img))
    images.append(img_tensor)

batched_images = torch.stack(images)  # [B, 3, H, W]
print(f"Batch shape: {batched_images.shape}")

# Define prompts for each image (same prompt for all in this example)
# Shape: [batch_size, num_queries, num_points, 2]
batched_points = torch.tensor([
    [[[256, 256]]],  # Image 1
    [[[256, 256]]],  # Image 2
    [[[256, 256]]],  # Image 3
    [[[256, 256]]],  # Image 4
], dtype=torch.float32)

# Shape: [batch_size, num_queries, num_points]
batched_labels = torch.tensor([
    [[1]],
    [[1]],
    [[1]],
    [[1]],
])

# Run batch inference
with torch.no_grad():
    predicted_logits, predicted_iou = model(
        batched_images,
        batched_points,
        batched_labels,
    )

print(f"Output logits shape: {predicted_logits.shape}")  # [B, num_queries, num_masks, H, W]
print(f"Output IoU shape: {predicted_iou.shape}")  # [B, num_queries, num_masks]

# Extract masks for each image
for i in range(batch_size):
    mask = torch.ge(predicted_logits[i, 0, 0], 0).cpu().numpy()
    iou = predicted_iou[i, 0, 0].item()
    print(f"Image {i+1}: IoU = {iou:.3f}, Mask pixels = {mask.sum()}")
```

--------------------------------

### Visualizing SqueezeSAM Segmentation

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Displays an image, runs SqueezeSAM for segmentation using provided points and labels, and then shows the resulting segmentation masks on the image. Requires 'image', 'input_point', 'input_label', 'run_ours_box_or_points', and 'show_anns_ours' to be defined.

```python
ax[3].imshow(image)
mask_squeeze_sam = run_ours_box_or_points(image_path, input_point, input_label, squeeze_sam_model)
show_anns_ours(mask_squeeze_sam, ax[3])
ax[3].title.set_text("SqueezeSAM")
ax[3].axis('off')
```

--------------------------------

### Import necessary libraries

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

Imports essential libraries for image processing, numerical operations, and deep learning model interaction. GRID_SIZE is a configurable parameter for the density of point prompts.

```python
import matplotlib.pyplot as plt
import numpy as np
import torch
from torchvision.transforms import ToTensor
from PIL import Image
import io
import cv2
GRID_SIZE = 32
```

--------------------------------

### Load EfficientSAM Model

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

Extracts model weights and initializes the EfficientSAM ViT-S model.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vits
import zipfile

with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
    zip_ref.extractall("weights")
efficient_sam_vits_model = build_efficient_sam_vits()
efficient_sam_vits_model.eval()
```

--------------------------------

### Decode with Different Prompts

Source: https://context7.com/yformer/efficientsam/llms.txt

This snippet demonstrates how to use the decoder session with different point and label queries to generate masks and predict IoU scores. Ensure image embeddings and the decoder session are pre-initialized.

```python
queries = [
    {"points": np.array([[[[580, 350], [650, 350]]]], dtype=np.float32),
     "labels": np.array([[[1, 1]]], dtype=np.float32)},
    {"points": np.array([[[[400, 200], [800, 600]]]], dtype=np.float32),
     "labels": np.array([[[2, 3]]], dtype=np.float32)},
]

for i, query in enumerate(queries):
    masks, iou_pred, _ = decoder_session.run(
        output_names=None,
        input_feed={
            "image_embeddings": image_embeddings,
            "batched_point_coords": query["points"],
            "batched_point_labels": query["labels"],
            "orig_im_size": np.array(input_image.shape[2:], dtype=np.int64),
        },
    )
    mask = masks[0, 0, 0] >= 0
    print(f"Query {i+1}: Mask shape = {mask.shape}, IoU = {iou_pred[0, 0, 0]:.3f}")
```

--------------------------------

### Import Necessary Libraries

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Imports essential libraries for image processing, numerical operations, and deep learning with PyTorch and Matplotlib.

```python
import matplotlib.pyplot as plt
import numpy as np
import torch
from torchvision.transforms import ToTensor
from PIL import Image
import io
```

--------------------------------

### Export EfficientSAM to TorchScript Format

Source: https://context7.com/yformer/efficientsam/llms.txt

This code exports EfficientSAM models (ViT and S variants) to TorchScript format for C++ deployment and optimized inference. It also shows how to load and run inference with a saved TorchScript model.

```python
import torch
import zipfile
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt, build_efficient_sam_vits

# Export EfficientSAM-Ti to TorchScript
model_ti = build_efficient_sam_vitt()
scripted_ti = torch.jit.script(model_ti)
torch.jit.save(scripted_ti, "efficient_sam_vitt_torchscript.pt")
print("Exported EfficientSAM-Ti to TorchScript")

# Export EfficientSAM-S to TorchScript
with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
    zip_ref.extractall("weights")
model_s = build_efficient_sam_vits()
scripted_s = torch.jit.script(model_s)
torch.jit.save(scripted_s, "efficient_sam_vits_torchscript.pt")
print("Exported EfficientSAM-S to TorchScript")

# Load and use TorchScript model
loaded_model = torch.jit.load("efficient_sam_vitt_torchscript.pt")
loaded_model.eval()

# Run inference with loaded model
image_tensor = torch.randn(1, 3, 512, 512)
points = torch.tensor([[[[256, 256]]]], dtype=torch.float32)
labels = torch.tensor([[[1]]])

with torch.no_grad():
    masks, iou = loaded_model(image_tensor, points, labels)

print(f"TorchScript output shape: {masks.shape}")
```

--------------------------------

### Point-Based Segmentation with EfficientSAM

Source: https://context7.com/yformer/efficientsam/llms.txt

Performs image segmentation using point prompts. Input points and labels define foreground objects. The output includes predicted masks sorted by IoU score.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt
from PIL import Image
from torchvision import transforms
import torch
import numpy as np

# Load model and image
model = build_efficient_sam_vitt()
model.eval()

image = Image.open("figs/examples/dogs.jpg")
image_np = np.array(image)
image_tensor = transforms.ToTensor()(image_np)

# Define point prompts - two foreground points on a dog
# Shape: [batch_size, num_queries, num_points, 2]
input_points = torch.tensor([[[[580, 350], [650, 350]]]])  # (x, y) coordinates
# Shape: [batch_size, num_queries, num_points]
input_labels = torch.tensor([[[1, 1]]])  # 1 = foreground point

# Run inference
with torch.no_grad():
    predicted_logits, predicted_iou = model(
        image_tensor[None, ...],  # Add batch dimension: [1, 3, H, W]
        input_points,
        input_labels,
    )

# Sort masks by predicted IoU score (highest quality first)
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)

# Extract best mask (threshold at 0)
# Output shape: [batch, num_queries, num_masks, H, W]
mask = torch.ge(predicted_logits[0, 0, 0, :, :], 0).cpu().numpy()

# Apply mask to image
masked_image = image_np.copy().astype(np.uint8) * mask[:, :, None]
Image.fromarray(masked_image).save("output_point_segmentation.png")

print(f"Mask shape: {mask.shape}")
print(f"IoU scores: {predicted_iou[0, 0].tolist()}")
```

--------------------------------

### Run Segmentation with Box or Points Input

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Processes an image using provided points and labels to generate segmentation masks. It sorts predictions by IoU and returns binary masks.

```python
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()
```

--------------------------------

### Run EfficientSAM Inference

Source: https://context7.com/yformer/efficientsam/llms.txt

Perform mask prediction and visualize results using PyTorch.

```python
# Run inference
with torch.no_grad():
    predicted_logits, predicted_iou = model(
        image_tensor[None, ...],
        input_points,
        input_labels,
    )

# Get best mask
sorted_ids = torch.argsort(predicted_iou, dim=-1, descending=True)
predicted_logits = torch.take_along_dim(predicted_logits, sorted_ids[..., None, None], dim=2)
mask = torch.ge(predicted_logits[0, 0, 0, :, :], 0).cpu().numpy()

# Visualize result
import matplotlib.pyplot as plt

fig, ax = plt.subplots(1, 2, figsize=(12, 6))
ax[0].imshow(image_np)
ax[0].add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1,
                               edgecolor='yellow', facecolor='none', lw=3))
ax[0].set_title("Input with Box Prompt")

ax[1].imshow(image_np)
mask_overlay = np.zeros((*mask.shape, 4))
mask_overlay[mask] = [0, 1, 0, 0.5]  # Green overlay
ax[1].imshow(mask_overlay)
ax[1].set_title("Segmentation Result")

plt.savefig("output_box_segmentation.png")
```

--------------------------------

### Box-Based Segmentation with EfficientSAM

Source: https://context7.com/yformer/efficientsam/llms.txt

Segments objects using a bounding box prompt. The input points define the top-left (label 2) and bottom-right (label 3) corners of the bounding box.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vits
from PIL import Image
from torchvision.transforms import ToTensor
import torch
import numpy as np
import zipfile

# Load model
with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
    zip_ref.extractall("weights")
model = build_efficient_sam_vits()
model.eval()

# Load image
image_path = "figs/examples/dogs.jpg"
image_np = np.array(Image.open(image_path))
image_tensor = ToTensor()(image_np)

# Define box prompt: top-left and bottom-right corners
x1, y1, x2, y2 = 400, 200, 800, 600  # Bounding box coordinates
input_points = torch.tensor([[[[x1, y1], [x2, y2]]]], dtype=torch.float32)
input_labels = torch.tensor([[[2, 3]]])  # 2 = top-left, 3 = bottom-right
```

--------------------------------

### Perform Inference with Point Prompts

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

Processes input points to generate segmentation masks using the EfficientSAM model.

```python
    xy = torch.from_numpy(np.array(xy))
    points = xy
    num_pts = xy.shape[0]
    point_labels = torch.ones(num_pts, 1)
    with torch.no_grad():
      predicted_masks, predicted_iou = get_predictions_given_embeddings_and_queries(
              img_tensor.cpu(),
              points.reshape(1, num_pts, 1, 2).cpu(),
              point_labels.reshape(1, num_pts, 1).cpu(),
              model.cpu(),
          )
    rle = [mask_to_rle_pytorch(m[0:1]) for m in predicted_masks]
    predicted_masks = process_small_region(rle)
    return predicted_masks
```

--------------------------------

### Run segmentation process for an image

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

This function orchestrates the segmentation process for a given image path. It loads the image, converts it to a tensor, and generates a grid of point prompts across the image.

```python
def run_everything_ours(img_path, model):
    model = model.cpu()
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    img_tensor = ToTensor()(image)
    _, original_image_h, original_image_w = img_tensor.shape
    xy = []
    for i in range(GRID_SIZE):
        curr_x = 0.5 + i / GRID_SIZE * original_image_w
        for j in range(GRID_SIZE):
            curr_y = 0.5 + j / GRID_SIZE * original_image_h
            xy.append([curr_x, curr_y])

```

--------------------------------

### EfficientSAM Attention and Projection Layers

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

This snippet shows the configuration of attention mechanisms and linear projection layers within the EfficientSAM model. It details the in-features and out-features for query, key, value, and output projections in the attention block, as well as normalization layers.

```python
      (final_attn_token_to_image): AttentionForTwoWayAttentionBlock(
        (q_proj): Linear(in_features=256, out_features=128, bias=True)
        (k_proj): Linear(in_features=256, out_features=128, bias=True)
        (v_proj): Linear(in_features=256, out_features=128, bias=True)
        (out_proj): Linear(in_features=128, out_features=256, bias=True)
      )
      (norm_final_attn): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
    )
    (iou_token): Embedding(1, 256)
    (mask_tokens): Embedding(4, 256)
    (final_output_upscaling_layers): ModuleList(
      (0): Sequential(
        (0): ConvTranspose2d(256, 64, kernel_size=(2, 2), stride=(2, 2))
        (1): GroupNorm(1, 64, eps=1e-05, affine=True)
        (2): GELU(approximate='none')
      )
      (1): Sequential(
        (0): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))
        (1): Identity()
        (2): GELU(approximate='none')
      )
    )
    (output_hypernetworks_mlps): ModuleList(
      (0-3): 4 x MLPBlock(
        (layers): ModuleList(
          (0-1): 2 x Sequential(
            (0): Linear(in_features=256, out_features=256, bias=True)
            (1): GELU(approximate='none')
          )
        )
        (fc): Linear(in_features=256, out_features=32, bias=True)
      )
    )
    (iou_prediction_head): MLPBlock(
      (layers): ModuleList(
        (0-1): 2 x Sequential(
          (0): Linear(in_features=256, out_features=256, bias=True)
          (1): GELU(approximate='none')
        )
      )
      (fc): Linear(in_features=256, out_features=4, bias=True)
    )
  )
)
```

--------------------------------

### Helper Functions for Visualization

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Provides utility functions to display segmentation masks, input points, and bounding boxes on a Matplotlib axes object.

```python
def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.8])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(
        pos_points[:, 0],
        pos_points[:, 1],
        color="green",
        marker="*",
        s=marker_size,
        edgecolor="white",
        linewidth=1.25,
    )
    ax.scatter(
        neg_points[:, 0],
        neg_points[:, 1],
        color="red",
        marker="*",
        s=marker_size,
        edgecolor="white",
        linewidth=1.25,
    )

def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(
        plt.Rectangle((x0, y0), w, h, edgecolor="yellow", facecolor=(0, 0, 0, 0), lw=5)
    )

def show_anns_ours(mask, ax):
    ax.set_autoscale_on(False)
    img = np.ones((mask.shape[0], mask.shape[1], 4))
    img[:, :, 3] = 0
    color_mask = [0, 1, 0, 0.7]
    img[np.logical_not(mask)] = color_mask
    ax.imshow(img)
```

--------------------------------

### Run ONNX Model Inference

Source: https://context7.com/yformer/efficientsam/llms.txt

Execute inference using ONNX Runtime with either a combined model or split encoder/decoder components.

```python
import numpy as np
import onnxruntime
from PIL import Image

# Load image and preprocess
image = np.array(Image.open("figs/examples/dogs.jpg"))
input_image = image.transpose(2, 0, 1)[None].astype(np.float32) / 255.0

# Define prompts
# Shape: [batch_size, num_queries, num_points, 2]
input_points = np.array([[[[580, 350], [650, 350]]]], dtype=np.float32)
# Shape: [batch_size, num_queries, num_points]
input_labels = np.array([[[1, 1]]], dtype=np.float32)

# Option 1: Single combined ONNX model
session = onnxruntime.InferenceSession("weights/efficient_sam_vitt.onnx")
output_masks, iou_predictions = session.run(
    output_names=None,
    input_feed={
        "batched_images": input_image,
        "batched_point_coords": input_points,
        "batched_point_labels": input_labels,
    },
)

print(f"Output masks shape: {output_masks.shape}")
print(f"IoU predictions: {iou_predictions}")

# Option 2: Separate encoder/decoder for multi-query efficiency
encoder_session = onnxruntime.InferenceSession("weights/efficient_sam_vitt_encoder.onnx")
decoder_session = onnxruntime.InferenceSession("weights/efficient_sam_vitt_decoder.onnx")

# Encode image once
image_embeddings, = encoder_session.run(
    output_names=None,
    input_feed={"batched_images": input_image},
)
```

--------------------------------

### Compute Image Embeddings for Multiple Queries

Source: https://context7.com/yformer/efficientsam/llms.txt

Cache image embeddings to process multiple prompts efficiently on the same image.

```python
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt
from PIL import Image
from torchvision import transforms
import torch
import numpy as np

model = build_efficient_sam_vitt()
model.eval()

# Load and preprocess image
image = Image.open("figs/examples/dogs.jpg")
image_np = np.array(image)
image_tensor = transforms.ToTensor()(image_np)[None, ...]  # [1, 3, H, W]

# Compute image embeddings once
with torch.no_grad():
    image_embeddings = model.get_image_embeddings(image_tensor)

print(f"Image embeddings shape: {image_embeddings.shape}")  # [1, 256, 64, 64]

# Define multiple query prompts for different objects
queries = [
    {"points": [[580, 350]], "labels": [1]},  # First dog
    {"points": [[200, 300]], "labels": [1]},  # Second object
    {"points": [[400, 200], [800, 600]], "labels": [2, 3]},  # Box query
]

batch_size, _, input_h, input_w = image_tensor.shape

# Process each query using cached embeddings
for i, query in enumerate(queries):
    points = torch.tensor([[query["points"]]], dtype=torch.float32)
    labels = torch.tensor([[query["labels"]]])

    with torch.no_grad():
        predicted_logits, predicted_iou = model.predict_masks(
            image_embeddings,
            points,
            labels,
            multimask_output=True,
            input_h=input_h,
            input_w=input_w,
            output_h=input_h,
            output_w=input_w,
        )

    mask = torch.ge(predicted_logits[0, 0, 0], 0).cpu().numpy()
    print(f"Query {i+1}: IoU = {predicted_iou[0, 0, 0].item():.3f}, Mask pixels = {mask.sum()}")
```

--------------------------------

### Export EfficientSAM to ONNX Format

Source: https://context7.com/yformer/efficientsam/llms.txt

This script converts EfficientSAM models to ONNX format, supporting dynamic axes for variable batch sizes and image dimensions. It handles full models, encoders, and decoders separately.

```python
import torch
import onnxruntime
from efficient_sam.build_efficient_sam import build_efficient_sam_vitt
from onnx_models import OnnxEfficientSam, OnnxEfficientSamEncoder, OnnxEfficientSamDecoder

def export_model(model, output_path, model_type="full"):
    """Export EfficientSAM to ONNX format."""

    if model_type == "full":
        onnx_model = OnnxEfficientSam(model=model)
        dynamic_axes = {
            "batched_images": {0: "batch", 2: "height", 3: "width"},
            "batched_point_coords": {2: "num_points"},
            "batched_point_labels": {2: "num_points"},
        }
        dummy_inputs = {
            "batched_images": torch.randn(1, 3, 1080, 1920, dtype=torch.float),
            "batched_point_coords": torch.randint(0, 1080, (1, 1, 5, 2), dtype=torch.float),
            "batched_point_labels": torch.randint(0, 4, (1, 1, 5), dtype=torch.float),
        }
        output_names = ["output_masks", "iou_predictions"]

    elif model_type == "encoder":
        onnx_model = OnnxEfficientSamEncoder(model=model)
        dynamic_axes = {"batched_images": {0: "batch", 2: "height", 3: "width"}}
        dummy_inputs = {"batched_images": torch.randn(1, 3, 1080, 1920, dtype=torch.float)}
        output_names = ["image_embeddings"]

    elif model_type == "decoder":
        onnx_model = OnnxEfficientSamDecoder(model=model)
        dynamic_axes = {
            "image_embeddings": {0: "batch"},
            "batched_point_coords": {2: "num_points"},
            "batched_point_labels": {2: "num_points"},
        }
        dummy_inputs = {
            "image_embeddings": torch.randn(1, 256, 64, 64, dtype=torch.float),
            "batched_point_coords": torch.randint(0, 1080, (1, 1, 5, 2), dtype=torch.float),
            "batched_point_labels": torch.randint(0, 4, (1, 1, 5), dtype=torch.float),
            "orig_im_size": torch.tensor([1080, 1920], dtype=torch.long),
        }
        output_names = ["output_masks", "iou_predictions"]

    torch.onnx.export(
        onnx_model,
        tuple(dummy_inputs.values()),
        output_path,
        export_params=True,
        opset_version=17,
        do_constant_folding=True,
        input_names=list(dummy_inputs.keys()),
        output_names=output_names,
        dynamic_axes=dynamic_axes,
    )
    print(f"Exported {model_type} model to {output_path}")

# Export all variants
model = build_efficient_sam_vitt()
export_model(model, "efficient_sam_vitt.onnx", "full")
export_model(model, "efficient_sam_vitt_encoder.onnx", "encoder")
export_model(model, "efficient_sam_vitt_decoder.onnx", "decoder")
```

--------------------------------

### Post-processing script for segmentation masks

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_segment_everything_example.ipynb

This script includes utility functions for post-processing segmentation masks, such as removing small regions and applying Non-Maximum Suppression (NMS). It's adapted from the original SAM project.

```python
from segment_anything.utils.amg import (
    batched_mask_to_box,
    calculate_stability_score,
    mask_to_rle_pytorch,
    remove_small_regions,
    rle_to_mask,
)
from torchvision.ops.boxes import batched_nms, box_area
def process_small_region(rles):
        new_masks = []
        scores = []
        min_area = 100
        nms_thresh = 0.7
        for rle in rles:
            mask = rle_to_mask(rle[0])

            mask, changed = remove_small_regions(mask, min_area, mode="holes")
            unchanged = not changed
            mask, changed = remove_small_regions(mask, min_area, mode="islands")
            unchanged = unchanged and not changed

            new_masks.append(torch.as_tensor(mask).unsqueeze(0))
            # Give score=0 to changed masks and score=1 to unchanged masks
            # so NMS will prefer ones that didn't need postprocessing
            scores.append(float(unchanged))

        # Recalculate boxes and remove any new duplicates
        masks = torch.cat(new_masks, dim=0)
        boxes = batched_mask_to_box(masks)
        keep_by_nms = batched_nms(
            boxes.float(),
            torch.as_tensor(scores),
            torch.zeros_like(boxes[:, 0]),  # categories
            iou_threshold=nms_thresh,
        )

        # Only recalculate RLEs for masks that have changed
        for i_mask in keep_by_nms:
            if scores[i_mask] == 0.0:
                mask_torch = masks[i_mask].unsqueeze(0)
                rles[i_mask] = mask_to_rle_pytorch(mask_torch)
        masks = [rle_to_mask(rles[i][0]) for i in keep_by_nms]
        return masks
```

--------------------------------

### EfficientSAM Model Architecture Details

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

This snippet shows the internal structure of the EfficientSAM model, including upsampling layers and convolutional blocks with activation and normalization. It details the layer dimensions and types used in the model's architecture.

```python
              (conv): Sequential(
                (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                (1): ReLU()
              )
              (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
              (quant_output_normalization): QuantStub()
            )
          )
        )
      )
      (1): Up(
        (up): ConvTranspose2dUsingLinear(
          (upsampler_linear): Sequential(
            (0): Linear(in_features=256, out_features=512, bias=True)
            (1): ReLU()
          )
          (norm): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (conv): DoubleConv(
          (double_conv): Sequential(
            (0): Conv2dWithActivation(
              (conv): Sequential(
                (0): Conv2d(292, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                (1): ReLU()
              )
              (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
              (quant_output_normalization): QuantStub()
            )
            (1): Conv2dWithActivation(
              (conv): Sequential(
                (0): Conv2d(164, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                (1): ReLU()
              )
              (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
              (quant_output_normalization): QuantStub()
            )
          )
        )
      )
      (2): Up(
        (up): ConvTranspose2dUsingLinear(
          (upsampler_linear): Sequential(
            (0): Linear(in_features=164, out_features=328, bias=True)
            (1): ReLU()
          )
          (norm): BatchNorm2d(328, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (conv): DoubleConv(
          (double_conv): Sequential(
            (0): Conv2dWithActivation(
              (conv): Sequential(
                (0): Conv2d(210, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                (1): ReLU()
              )
              (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
              (quant_output_normalization): QuantStub()
            )
            (1): Conv2dWithActivation(
              (conv): Sequential(
                (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                (1): ReLU()
              )
              (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
              (quant_output_normalization): QuantStub()
            )
          )
        )
      )
      (3-5): 3 x Up(
        (up): ConvTranspose2dUsingLinear(
          (upsampler_linear): Sequential(
            (0): Linear(in_features=128, out_features=256, bias=True)
            (1): ReLU()
          )
          (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (conv): DoubleConv(
          (double_conv): Sequential(

```

--------------------------------

### SqueezeSam Model Architecture

Source: https://github.com/yformer/efficientsam/blob/main/notebooks/EfficientSAM_example.ipynb

Defines the SqueezeSam model architecture, including its UnetEncoder and various convolutional and normalization layers.

```python
SqueezeSam(
  (image_encoder): UnetEncoder(
    (proj): Sequential(
      (0): Conv2d(5, 64, kernel_size=(2, 2), stride=(2, 2))
      (1): GroupNorm(1, 64, eps=1e-05, affine=True)
    )
    (down_layers): ModuleList(
      (0): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2dWithActivation(
            (conv): Sequential(
              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
              (1): ReLU()
            )
            (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            (quant_output_normalization): QuantStub()
          )
          (1): Conv2dWithActivation(
            (conv): Sequential(
              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
              (1): ReLU()
            )
            (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            (quant_output_normalization): QuantStub()
          )
        )
      )
      (1): Down(
        (maxpool_conv): Sequential(
          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (1): DoubleConv(
            (double_conv): Sequential(
              (0): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )
                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                (quant_output_normalization): QuantStub()
              )
              (1): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )
                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                (quant_output_normalization): QuantStub()
              )
            )
          )
        )
      )
      (2-4): 3 x Down(
        (maxpool_conv): Sequential(
          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (1): DoubleConv(
            (double_conv): Sequential(
              (0): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )
                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                (quant_output_normalization): QuantStub()
              )
              (1): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )
                (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                (quant_output_normalization): QuantStub()
              )
            )
          )
        )
      )
      (5): Down(
        (maxpool_conv): Sequential(
          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (1): DoubleConv(
            (double_conv): Sequential(
              (0): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(128, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )
                (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                (quant_output_normalization): QuantStub()
              )
              (1): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(164, 164, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )
                (norm): BatchNorm2d(164, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                (quant_output_normalization): QuantStub()
              )
            )
          )
        )
      )
      (6): Down(
        (maxpool_conv): Sequential(
          (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (1): DoubleConv(
            (double_conv): Sequential(
              (0): Conv2dWithActivation(
                (conv): Sequential(
                  (0): Conv2d(164, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                  (1): ReLU()
                )

```

--------------------------------

### Cite EfficientSAM

Source: https://github.com/yformer/efficientsam/blob/main/README.md

BibTeX citation for research or application usage.

```bibtex
@article{xiong2023efficientsam,
  title={EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything},
  author={Yunyang Xiong, Bala Varadarajan, Lemeng Wu, Xiaoyu Xiang, Fanyi Xiao, Chenchen Zhu, Xiaoliang Dai, Dilin Wang, Fei Sun, Forrest Iandola, Raghuraman Krishnamoorthi, Vikas Chandra},
  journal={arXiv:2312.00863},
  year={2023}
}
```
Tokens:
Raw