Image Prompter: Segment Anything

This tutorials shows how to use our high-level API Image Prompter. This API allow to set an image, and run multiple queries multiple times on this image. These query can be done with three types of prompt:

1. Points: Keypoints with (x, y) and a respective label. Where 0 indicates a background point; 1 indicates a foreground point;

2. Boxes: Boxes of different regions.

3. Masks: Logits generated by the model in a previous run.

Read more on our docs: https://kornia.readthedocs.io/en/latest/models/segment_anything.html

This tutorials steps:

1. Setup the desired SAM model and import the necessary packages

   1. Utilities function to read and plot the data

2. How to instantiate the Image Prompter

   1. How to set an image

   2. The supported prompts type

   3. Example how to query on the image

   4. Prediction structure

3. Using the Image Prompter

   1. Soccer player segmentation

   2. Car parts segmentation

   3. Satellite image - Sentinel-2

Setup

%%capture
!pip install kornia
!pip install kornia-rs

First let’s choose the SAM type to be used on our Image Prompter.

The options are (smaller to bigger):

model_type	checkpoint official
vit_b	https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
vit_l	https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth
vit_h	https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

model_type = "vit_h"
checkpoint = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"

Then let’s import all necessary packages and modules

from __future__ import annotations

import os

import matplotlib.pyplot as plt
import torch
from kornia.contrib.image_prompter import ImagePrompter
from kornia.contrib.models.sam import SamConfig
from kornia.geometry.boxes import Boxes
from kornia.geometry.keypoints import Keypoints
from kornia.io import ImageLoadType, load_image
from kornia.utils import get_cuda_device_if_available, tensor_to_image

device = get_cuda_device_if_available()
print(device)

cuda:0

Utilities functions

import io

import requests


def download_image(url: str, filename: str = "") -> str:
    filename = url.split("/")[-1] if len(filename) == 0 else filename
    # Download
    bytesio = io.BytesIO(requests.get(url).content)
    # Save file
    with open(filename, "wb") as outfile:
        outfile.write(bytesio.getbuffer())

    return filename


soccer_image_path = download_image("https://raw.githubusercontent.com/kornia/data/main/soccer.jpg")
car_image_path = download_image("https://raw.githubusercontent.com/kornia/data/main/simple_car.jpg")
satellite_image_path = download_image("https://raw.githubusercontent.com/kornia/data/main/satellite_sentinel2_example.tif")
soccer_image_path, car_image_path, satellite_image_path

('soccer.jpg', 'simple_car.jpg', 'satellite_sentinel2_example.tif')

def colorize_masks(binary_masks: torch.Tensor, merge: bool = True, alpha: None | float = None) -> list[torch.Tensor]:
    """Convert binary masks (B, C, H, W), boolean tensors, into masks with colors (B, (3, 4) , H, W) - RGB or RGBA. Where C refers to the number of masks.
    Args:
        binary_masks: a batched boolean tensor (B, C, H, W)
        merge: If true, will join the batch dimension into a unique mask.
        alpha: alpha channel value. If None, will generate RGB images

    Returns:
        A list of `C` colored masks.
    """
    B, C, H, W = binary_masks.shape
    OUT_C = 4 if alpha else 3

    output_masks = []

    for idx in range(C):
        _out = torch.zeros(B, OUT_C, H, W, device=binary_masks.device, dtype=torch.float32)
        for b in range(B):
            color = torch.rand(1, 3, 1, 1, device=binary_masks.device, dtype=torch.float32)
            if alpha:
                color = torch.cat([color, torch.tensor([[[[alpha]]]], device=binary_masks.device, dtype=torch.float32)], dim=1)

            to_colorize = binary_masks[b, idx, ...].view(1, 1, H, W).repeat(1, OUT_C, 1, 1)
            _out[b, ...] = torch.where(to_colorize, color, _out[b, ...])
        output_masks.append(_out)

    if merge:
        output_masks = [c.max(dim=0)[0] for c in output_masks]

    return output_masks


def show_binary_masks(binary_masks: torch.Tensor, axes) -> None:
    """plot binary masks, with shape (B, C, H, W), where C refers to the number of masks.

    will merge the `B` channel into a unique mask.
    Args:
        binary_masks: a batched boolean tensor (B, C, H, W)
        ax: a list of matplotlib axes with lenght of C
    """
    colored_masks = colorize_masks(binary_masks, True, 0.6)

    for ax, mask in zip(axes, colored_masks):
        ax.imshow(tensor_to_image(mask))


def show_boxes(boxes: Boxes, ax) -> None:
    boxes_tensor = boxes.to_tensor(mode="xywh").detach().cpu().numpy()
    for batched_boxes in boxes_tensor:
        for box in batched_boxes:
            x0, y0, w, h = box
            ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor="orange", facecolor=(0, 0, 0, 0), lw=2))


def show_points(points: tuple[Keypoints, torch.Tensor], ax, marker_size=200):
    coords, labels = points
    pos_points = coords[labels == 1].to_tensor().detach().cpu().numpy()
    neg_points = coords[labels == 0].to_tensor().detach().cpu().numpy()

    ax.scatter(pos_points[:, 0], pos_points[:, 1], color="green", marker="+", s=marker_size, linewidth=2)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color="red", marker="x", s=marker_size, linewidth=2)

from kornia.contrib.models import SegmentationResults


def show_image(image: torch.Tensor):
    plt.imshow(tensor_to_image(image))
    plt.axis("off")
    plt.show()


def show_predictions(
    image: torch.Tensor,
    predictions: SegmentationResults,
    points: tuple[Keypoints, torch.Tensor] | None = None,
    boxes: Boxes | None = None,
) -> None:
    n_masks = predictions.logits.shape[1]

    fig, axes = plt.subplots(1, n_masks, figsize=(21, 16))
    axes = [axes] if n_masks == 1 else axes

    for idx, ax in enumerate(axes):
        score = predictions.scores[:, idx, ...].mean()
        ax.imshow(tensor_to_image(image))
        ax.set_title(f"Mask {idx+1}, Score: {score:.3f}", fontsize=18)

        if points:
            show_points(points, ax)

        if boxes:
            show_boxes(boxes, ax)

        ax.axis("off")

    show_binary_masks(predictions.binary_masks, axes)
    plt.show()

Exploring the Image Prompter

The ImagePrompter can be initialized from a ModelConfig structure, where now we just have support for the SAM model through the SamConfig. Through this config the ImagePrompter will initialize the SAM model and load the weights (from a path or a URL).

What the ImagePrompter can do?

1. Based on the ModelConfig, besides the model initialization, we will setup the required transformations for the images and prompts using the kornia.augmentation API within the Augmentation sequential container.

2. You can benefit from using the torch.compile(...) API (dynamo) for torch >= 2.0.0 versions. To compile with dynamo we provide the method ImagePrompter.compile(...) which will optimize the right parts of the backend model and the prompter itself.

3. Caching the image features and transformations. With the ImagePrompter.set_image(...) method, we transform the image and already encode it using the model, caching it’s embeddings to query later.

4. Query multiple times with multiple prompts. Using the ImagePrompter.predict(...), where we will query on our cached embeddings using Keypoints, Boxes and Masks as prompt.

What the ImagePrompter and Kornia provides? Easy high-levels structures to be used as prompt, also as the result of the prediction. Using the kornia geometry module you can easily encapsulate the Keypoints and Boxes, which allow the API to be more flexible about the desired mode (mainly for boxes, where we had multiple modes of represent it).

The Kornia ImagePrompter and model config for SAM can be imported as follow:

from kornia.contrib.image_prompter import ImagePrompter
from kornia.contrib.models import SamConfig

# Setting up a SamConfig with the model type and checkpoint desired
config = SamConfig(model_type, checkpoint)

# Initialize the ImagePrompter
prompter = ImagePrompter(config, device=device)

Set image

First, before set the image into the prompter, we need to read the image. For it, we can use kornia.io, which internally uses kornia-rs. So for, it ensure to have kornia-rs installed, you can install it with pip install kornia_rs. This API implement the DLPack protocol natively in Rust to reduce the memory footprint during the decoding and types conversion. Allowing us to read the image from the disk directly to a tensor.

# Load the image
image = load_image(soccer_image_path, ImageLoadType.RGB32, device)  # 3 x H x W

# Display the loaded image
show_image(image)

With the image loaded into the same device than the model, and with the right shape 3xHxW let’s set the image into our image prompter. Attention, when doing this the model will already compute the embeddings of this image. This means, we will pass this image through the encoder, which will uses a lot of memory. It is possible to use the largest model (vit-h) with a graphic card (GPU) that has at least 8Gb of VRAM.

prompter.set_image(image)

If no error occurred, the features needed to run queries are already cached. If you want to check this, you can see the status of the prompter.is_image_set property.

prompter.is_image_set

True

Examples of prompts

The ImagePrompter output will have the same Batch Size that its prompts. Where the output shape will be (B, C, H, W). Where B is the number of input prompts, C is determined by multimask output parameter. If multimask_output is True than C=3, otherwise C=1

Keypoints

Keypoints prompts is a tensor or a Keypoint structure within coordinates into (x, y). With shape BxNx2.

For each coordinate pair, should have a corresponding label, where 0 indicates a background point; 1 indicates a foreground point; These labels should be in a tensor with shape BxN

The model will try to find a object within all the foreground points, and without the background points. In other words, the foreground points can be used to select the desired type of data, and the background point to exclude the type of data.

keypoints_tensor = torch.tensor([[[960, 540]]], device=device, dtype=torch.float32)
keypoints = Keypoints(keypoints_tensor)

labels = torch.tensor([[1]], device=device, dtype=torch.float32)

Boxes

Boxes prompts is a tensor a with boxes on “xyxy” format/mode, or a Boxes structure. Tensor should have a shape of BxNx4.

boxes_tensor = torch.tensor([[[1841.7, 739.0, 1906.5, 890.6]]], device=device, dtype=torch.float32)
boxes = Boxes.from_tensor(boxes_tensor, mode="xyxy")

Masks

Masks prompts should be provide from a previous model output, with shape Bx1x256x256

# first run
predictions = prompter.prediction(...)

# use previous results as prompt
predictions = prompter.prediction(..., mask=predictions.logits)

Example of prediction

# using keypoints
prediction_by_keypoint = prompter.predict(keypoints, labels, multimask_output=False)

show_image(prediction_by_keypoint.binary_masks)

# Using boxes
prediction_by_boxes = prompter.predict(boxes=boxes, multimask_output=False)

show_image(prediction_by_boxes.binary_masks)

Exploring the prediction result structure

The ImagePrompter prediction structure, is a SegmentationResults which has the upscaled (default) logits when output_original_size=True is passed on the predict.

The segmentation results have:

• logits: Results logits with shape (B, C, H, W), where C refers to the number of predicted masks

• scores: The scores from the logits. Shape (B,)

• Binary mask generated from logits considering the mask_threshold. The size depends on original_res_logits=True, if false, the binary mask will have the same shape of the logits Bx1x256x256

prediction_by_boxes.scores

tensor([[0.9317]], device='cuda:0')

prediction_by_boxes.binary_masks.shape

torch.Size([1, 1, 1080, 1920])

prediction_by_boxes.logits.shape

torch.Size([1, 1, 256, 256])

Using the Image Prompter on examples

Soccer players

Using an example image from the dataset: https://www.kaggle.com/datasets/ihelon/football-player-segmentation

Lets segment the persons on the field using boxes

# Prompts
boxes = Boxes.from_tensor(
    torch.tensor(
        [
            [
                [1841.7000, 739.0000, 1906.5000, 890.6000],
                [879.3000, 545.9000, 948.2000, 669.2000],
                [55.7000, 595.0000, 127.4000, 745.9000],
                [1005.4000, 128.7000, 1031.5000, 212.0000],
                [387.4000, 424.1000, 438.2000, 539.0000],
                [921.0000, 377.7000, 963.3000, 483.0000],
                [1213.2000, 885.8000, 1276.2000, 1060.1000],
                [40.8900, 725.9600, 105.4100, 886.5800],
                [848.9600, 283.6200, 896.0600, 368.6200],
                [1109.6500, 499.0400, 1153.0400, 622.1700],
                [576.3000, 860.8000, 671.7000, 1018.8000],
                [1039.8000, 389.9000, 1072.5000, 493.2000],
                [1647.1000, 315.1000, 1694.0000, 406.0000],
                [1231.2000, 214.0000, 1294.1000, 297.3000],
            ]
        ],
        device=device,
    ),
    mode="xyxy",
)

# Load the image
image = load_image(soccer_image_path, ImageLoadType.RGB32, device)  # 3 x H x W

# Set the image
prompter.set_image(image)

predictions = prompter.predict(boxes=boxes, multimask_output=True)

let’s see the results, since we used multimask_output=True, the model outputted 3 masks.

show_predictions(image, predictions, boxes=boxes)

Car parts

Segmenting car parts of an example from the dataset: https://www.kaggle.com/datasets/jessicali9530/stanford-cars-dataset

# Prompts
boxes = Boxes.from_tensor(
    torch.tensor(
        [
            [
                [56.2800, 369.1100, 187.3000, 579.4300],
                [412.5600, 426.5800, 592.9900, 608.1600],
                [609.0800, 366.8200, 682.6400, 431.1700],
                [925.1300, 366.8200, 959.6100, 423.1300],
                [756.1900, 416.2300, 904.4400, 473.7000],
                [489.5600, 285.2200, 676.8900, 343.8300],
            ]
        ],
        device=device,
    ),
    mode="xyxy",
)

keypoints = Keypoints(
    torch.tensor(
        [[[535.0, 227.0], [349.0, 215.0], [237.0, 219.0], [301.0, 373.0], [641.0, 397.0], [489.0, 513.0]]], device=device
    )
)
labels = torch.ones(keypoints.shape[:2], device=device, dtype=torch.float32)

# Image
image = load_image(car_image_path, ImageLoadType.RGB32, device)

# Set the image
prompter.set_image(image)

Querying with boxes

predictions = prompter.predict(boxes=boxes, multimask_output=True)

show_predictions(image, predictions, boxes=boxes)

Querying with keypoints

Considering N points into 1 Batch

This way the model will kinda find the object within all the points

predictions = prompter.predict(keypoints=keypoints, keypoints_labels=labels)

show_predictions(image, predictions, points=(keypoints, labels))

Considering 1 point into N Batch

Prompter encoder not working for a batch of points :/

Considering 1 point for batch into N queries

This way the model will find an object for each point

k = 2  # number of times/points to query

for idx in range(min(keypoints.data.size(1), k)):
    print("-" * 79, f"\nQuery {idx}:")
    _kpts = keypoints[:, idx, ...][None, ...]
    _lbl = labels[:, idx, ...][None, ...]

    predictions = prompter.predict(keypoints=_kpts, keypoints_labels=_lbl)

    show_predictions(image, predictions, points=(_kpts, _lbl))

------------------------------------------------------------------------------- 
Query 0:

------------------------------------------------------------------------------- 
Query 1:

Satellite image

Image from Sentinel-2

Product: A tile of the TCI (px of 10m). Product name: S2B_MSIL1C_20230324T130249_N0509_R095_T23KPQ_20230324T174312

# Prompts
keypoints = Keypoints(
    torch.tensor(
        [
            [
                # Urban
                [74.0, 104.5],
                [335, 110],
                [702, 65],
                [636, 479],
                [408, 820],
                # Forest
                [40, 425],
                [680, 566],
                [405, 439],
                [73, 689],
                [865, 460],
                # Ocean/water
                [981, 154],
                [705, 714],
                [357, 683],
                [259, 908],
                [1049, 510],
            ]
        ],
        device=device,
    )
)
labels = torch.ones(keypoints.shape[:2], device=device, dtype=torch.float32)

# Image
image = load_image(satellite_image_path, ImageLoadType.RGB32, device)

# Set the image
prompter.set_image(image)

Query urban points

# Query the prompts
labels_to_query = labels.clone()
labels_to_query[..., 5:] = 0

predictions = prompter.predict(keypoints=keypoints, keypoints_labels=labels_to_query)

show_predictions(image, predictions, points=(keypoints, labels_to_query))

Query Forest points

# Query the prompts
labels_to_query = labels.clone()
labels_to_query[..., :5] = 0
labels_to_query[..., 10:] = 0

predictions = prompter.predict(keypoints=keypoints, keypoints_labels=labels_to_query)

show_predictions(image, predictions, points=(keypoints, labels_to_query))

Query ocean/water points

# Query the prompts
labels_to_query = labels.clone()
labels_to_query[..., :10] = 0

predictions = prompter.predict(keypoints=keypoints, keypoints_labels=labels_to_query)

show_predictions(image, predictions, points=(keypoints, labels_to_query))