Implementation:Kornia Kornia Visual Prompter

Knowledge Sources	Kornia
Domains	Vision, Segmentation, Prompt_Based_Vision
Last Updated	2026-02-09 15:00 GMT

Overview

Provides a prompt-based visual segmentation interface wrapping the Segment Anything Model (SAM) for interactive mask prediction using points, boxes, and mask prompts.

Description

The VisualPrompter class in the Kornia contrib package wraps the SAM (Segment Anything Model) to enable interactive, prompt-based image segmentation. The user first sets an image via set_image, which resizes, normalizes, and encodes the image. Then, predict accepts point prompts (keypoints with labels), bounding box prompts, or low-resolution mask prompts to generate segmentation masks. The class handles all necessary transformations between the original image coordinate space and the encoder's input space, including prompt reprojection and output upscaling. It supports torch.compile for performance optimization.

Usage

Import this module when you need to perform interactive or prompt-based image segmentation using SAM, when you want to segment objects by specifying points, bounding boxes, or prior masks as prompts.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/contrib/visual_prompter.py
• Lines: 1-414

Signature

class VisualPrompter:
    def __init__(
        self,
        config: Optional[SamConfig] = None,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> None: ...

    def preprocess_image(
        self, x: torch.Tensor, mean: Optional[torch.Tensor] = None,
        std: Optional[torch.Tensor] = None
    ) -> torch.Tensor: ...

    def set_image(
        self, image: torch.Tensor, mean: Optional[torch.Tensor] = None,
        std: Optional[torch.Tensor] = None
    ) -> None: ...

    def predict(
        self,
        keypoints: Optional[Keypoints | torch.Tensor] = None,
        keypoints_labels: Optional[torch.Tensor] = None,
        boxes: Optional[Boxes | torch.Tensor] = None,
        masks: Optional[torch.Tensor] = None,
        multimask_output: bool = True,
        output_original_size: bool = True,
    ) -> SegmentationResults: ...

    def reset_image(self) -> None: ...

    def compile(self, *, fullgraph=False, dynamic=False,
                backend="inductor", mode=None, options=None, disable=False) -> None: ...

Import

from kornia.contrib import VisualPrompter
from kornia.models.sam import SamConfig

I/O Contract

Inputs (set_image)

Name	Type	Required	Description
image	torch.Tensor	Yes	RGB image tensor with shape (3, H, W), float32 in range [0, 1]
mean	torch.Tensor	No	Custom per-channel mean for normalization
std	torch.Tensor	No	Custom per-channel standard deviation for normalization

Inputs (predict)

Name	Type	Required	Description
keypoints	Keypoints or torch.Tensor	No	Point prompts with shape (K, N, 2) in pixel coordinates (X, Y)
keypoints_labels	torch.Tensor	No	Labels for point prompts with shape (K, N); 1=foreground, 0=background
boxes	Boxes or torch.Tensor	No	Bounding box prompts with shape (K, 4) in xyxy format
masks	torch.Tensor	No	Low-resolution mask prompts with shape (K, 1, 256, 256)
multimask_output	bool	No	If True, returns 3 masks per prompt (default: True)
output_original_size	bool	No	If True, upscales logits to original image size (default: True)

Outputs

Name	Type	Description
result	SegmentationResults	Contains logits (mask predictions) and scores (IoU confidence scores)

Usage Examples

import torch
from kornia.contrib import VisualPrompter
from kornia.models.sam import SamConfig
from kornia.geometry.boxes import Boxes

# Initialize with ViT-B SAM model
prompter = VisualPrompter(SamConfig('vit_b'))

# Set the image to segment
image = torch.rand(3, 512, 512)
prompter.set_image(image)

# Predict using a bounding box prompt
boxes = Boxes(
    torch.tensor([[[[100, 100], [100, 300], [300, 100], [300, 300]]]], dtype=torch.float32),
    mode='xyxy'
)
prediction = prompter.predict(boxes=boxes)
print(prediction.logits.shape)   # (1, 3, 256, 256) with multimask_output=True
print(prediction.scores.shape)   # (1, 3) IoU scores for each mask

# Reset to process a new image
prompter.reset_image()

Related Pages

• Principle:Kornia_Kornia_Visual_Prompting