DEV Community

Cover image for From a click to a mask: how point prompts drive SAM for object selection
Wlad Radchenko
Wlad Radchenko

Posted on

From a click to a mask: how point prompts drive SAM for object selection

You click once on a cat in a photo and a clean mask snaps to its outline. Click again on the background and the mask shrinks away from a spot you did not want. It feels instant, like the model is re-running the whole thing on every click. It is not. The expensive part ran once, when the image loaded. Every click after that is cheap.

This is a walkthrough of how that works in the Segment Anything code inside Wunjo Make. I will follow the real prompt path: encode the image once with set_image, then send each click as a tiny prompt to the decoder. The model behind it is SAM (Kirillov et al., ICCV 2023), and the split between a heavy image encoder and a light prompt decoder is the whole reason interactive selection feels live.

A glowing neural mesh wrapping a shape
The heavy work happens once. After that, each click is a quick lookup against what the model already knows about the picture. Photo: Unsplash.

The one expensive step, done once

When you load an image, the predictor runs it through the image encoder and keeps the result:

@torch.no_grad()
def set_torch_image(self, transformed_image, original_image_size):
    self.reset_image()
    self.original_size = original_image_size
    self.input_size = tuple(transformed_image.shape[-2:])
    input_image = self.model.preprocess(transformed_image)
    self.features = self.model.image_encoder(input_image)
    self.is_image_set = True
Enter fullscreen mode Exit fullscreen mode

That self.model.image_encoder(input_image) call is the costly one. It is a vision transformer chewing through a 1024x1024 image and producing a feature grid. The wrapper in Wunjo Make caches the output so you only pay for it once per image:

@staticmethod
def get_embedding(predictor, img):
    predictor.set_image(img)
    return predictor.get_image_embedding().cpu().numpy()
Enter fullscreen mode Exit fullscreen mode

The embedding is small and fixed. For SAM it is a 1 x 256 x 64 x 64 tensor, a 64x64 grid where each cell carries a 256-dimensional description of that patch of the picture. Once you have it, you never touch the encoder again for that image. Every click reads against this grid.

If you try to predict before this runs, the code stops you cold:

if not self.is_image_set:
    raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
Enter fullscreen mode Exit fullscreen mode

So the rule is simple. Set the image once. Then click as much as you want.

A click is two numbers and a label

In the browser the user clicks. The wrapper's draw_mask turns those clicks into two arrays the model understands: where the points are, and what each point means.

for point in point_list:
    # scale the click from canvas pixels to original image pixels
    scaleFactorX = originalWidth / canvasWidth
    scaleFactorY = originalHeight / canvasHeight
    newX = int(point['x'] * scaleFactorX)
    newY = int(point['y'] * scaleFactorY)
    input_point.append([newX, newY])

    # color tells us if this is a "keep" or "drop" click
    if point.get("color") == "lightblue" or point.get("color") == 1:
        input_label.append(1)   # foreground: include this
    elif point.get("color") == "red" or point.get("color") == 0:
        input_label.append(0)   # background: exclude this
Enter fullscreen mode Exit fullscreen mode

That label is the whole interaction model. A label of 1 is a foreground click, meaning "this is part of the thing I want." A label of 0 is a background click, meaning "this is not part of it, push the mask away from here." Blue clicks add, red clicks subtract. The user paints intent with two colors, and the model figures out the rest.

The first click is deliberately ambiguous. One blue dot on a person could mean the shirt, the torso, or the whole body. So you click again. Each extra click narrows it down. Background clicks are how you carve away the bits SAM grabbed that you did not want.

A box is just two more points with special labels

SAM treats a bounding box as two corner points, but with their own labels so the model knows they are corners and not regular clicks:

if box is not None:
    onnx_box_coords = box.reshape(2, 2)
    onnx_box_labels = np.array([2, 3])   # 2 = top-left corner, 3 = bottom-right corner
    onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]
    onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)
Enter fullscreen mode Exit fullscreen mode

So the label vocabulary is small and fixed: 1 foreground point, 0 background point, 2 box top-left, 3 box bottom-right. You can mix them. A box to rough out the region plus a couple of foreground and background clicks to clean up the edges is a strong combination.

When there is no box, the code still appends one filler point with label -1:

else:
    onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]
    onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)
Enter fullscreen mode Exit fullscreen mode

That -1 is a padding label, a "not a point" marker. It exists so the input tensor has a consistent shape even when you only sent clicks. Inside the prompt encoder, those padded slots get zeroed out and replaced with a dedicated embedding so they contribute nothing to the result. More on that next.

Where the labels turn into something the model can use

The clicks are still just coordinates and integers. The prompt encoder is where they become vectors the decoder can read. This is the heart of it:

def _embed_points(self, points, labels, pad):
    points = points + 0.5  # shift to the center of the pixel
    if pad:
        padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
        padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
        points = torch.cat([points, padding_point], dim=1)
        labels = torch.cat([labels, padding_label], dim=1)
    point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
    point_embedding[labels == -1] = 0.0
    point_embedding[labels == -1] += self.not_a_point_embed.weight
    point_embedding[labels == 0]  += self.point_embeddings[0].weight
    point_embedding[labels == 1]  += self.point_embeddings[1].weight
    return point_embedding
Enter fullscreen mode Exit fullscreen mode

Read it from the bottom. Each point first gets a position embedding from forward_with_coords, which encodes where the click landed. Then the code adds a learned vector that encodes what kind of click it was. Label 0 gets one learned vector, label 1 gets a different one, label -1 (the padding) gets wiped to zero and given the not_a_point vector so it carries no location signal. Boxes get the same treatment with their own two learned vectors:

def _embed_boxes(self, boxes):
    boxes = boxes + 0.5
    coords = boxes.reshape(-1, 2, 2)
    corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size)
    corner_embedding[:, 0, :] += self.point_embeddings[2].weight  # top-left
    corner_embedding[:, 1, :] += self.point_embeddings[3].weight  # bottom-right
    return corner_embedding
Enter fullscreen mode Exit fullscreen mode

There are exactly four learned point vectors, set up once in the constructor:

self.num_point_embeddings = 4  # pos/neg point + 2 box corners
point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)]
Enter fullscreen mode Exit fullscreen mode

Four small vectors. Foreground, background, box corner one, box corner two. That is the entire prompt language. Position says where, the learned vector says what role, and the two added together is what the decoder reads. The output of all this is a handful of tiny vectors, the "sparse embeddings." That is the cheap thing that flows into the decoder on every click.

The cheap step that runs on every click

Now the payoff. With the image embedding already cached, predicting a mask is just the prompt encoder plus the mask decoder:

@torch.no_grad()
def predict_torch(self, point_coords, point_labels, boxes, mask_input, multimask_output, return_logits=False):
    points = (point_coords, point_labels) if point_coords is not None else None

    # cheap: turn the clicks into a few vectors
    sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
        points=points, boxes=boxes, masks=mask_input,
    )

    # cheap: the decoder reads the cached image grid against those vectors
    low_res_masks, iou_predictions = self.model.mask_decoder(
        image_embeddings=self.features,                 # cached from set_image
        image_pe=self.model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,     # from your clicks
        dense_prompt_embeddings=dense_embeddings,
        multimask_output=multimask_output,
    )
    masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size)
    if not return_logits:
        masks = masks > self.model.mask_threshold
    return masks, iou_predictions, low_res_masks
Enter fullscreen mode Exit fullscreen mode

Notice image_embeddings=self.features. The decoder does not re-encode anything. It reads the grid the encoder produced once and lets the prompt vectors point at the part you mean. The decoder is small, the prompts are tiny, so this runs in a few milliseconds even on modest hardware. That is the entire reason clicking feels live. The cost of a click is the cost of the decoder, not the encoder.

Three masks, with scores, so ambiguity is solvable

One click is genuinely ambiguous, so by default SAM returns three candidate masks instead of guessing one:

multimask_output: bool = True
Enter fullscreen mode Exit fullscreen mode

The docstring in the code spells out why:

If true, the model will return three masks. For ambiguous input prompts (such as a single click), this will often produce better masks than a single prediction. If only a single mask is needed, the model's predicted quality score can be used to select the best mask.

Each candidate comes with a predicted quality score, the iou_predictions. So you can show all three and let the user pick, or auto-pick the highest-scoring one. The return tuple is exactly that:

masks_np = masks[0].detach().cpu().numpy()                 # the masks
iou_predictions_np = iou_predictions[0].detach().cpu().numpy()  # one score per mask
low_res_masks_np = low_res_masks[0].detach().cpu().numpy()      # 256x256 logits to reuse
Enter fullscreen mode Exit fullscreen mode

Once your prompts are unambiguous, say a box plus a few clicks, you can flip multimask_output to False and get a single cleaner mask. The code's own guidance: ambiguous input wants three, well-specified input wants one.


Wlad Radchenko About the author. I'm Wlad Radchenko, a software engineer. The code in this article comes from Wunjo Make (open source), local software for video makers, and Wunjo Design, an offline PWA for designers. Get in touch to find more on GitHub and LinkedIn.

Gotchas I would warn you about

A few things that bite people working with this path:

  • Set the image once, not per click. The easy mistake is calling set_image inside your click handler. That re-runs the encoder every time and throws away the entire speed advantage. Encode on load, predict on click. The wrapper's draw_mask does call set_image because it is built around per-frame masking for video; if you are masking one still image across many clicks, hoist the embedding out and reuse it.
  • Coordinates are in original image pixels, not canvas pixels. The frontend canvas is usually a scaled-down preview. If you skip the scaleFactorX/scaleFactorY step, your clicks land in the wrong place and the mask makes no sense. Always rescale before sending.
  • Points are (X, Y), and the labels array must match. Point coords are [x, y], and there must be exactly one label per point. point_labels must be supplied if point_coords is supplied is an assertion in the code, not a suggestion.
  • The -1 padding point is not optional with the ONNX path. It keeps the input shape consistent when there is no box. Drop it and the exported model gets the wrong tensor shape.
  • The returned low_res_masks is reusable. Those 256x256 logits can be fed back as mask_input on the next click to refine the same object. It is the cheapest way to iterate, and it is what makes "add one more click" feel additive rather than a fresh start.

Takeaway

The mechanism is a clean split. The image encoder is heavy and runs once, producing a small cached feature grid. Each click is turned into a few tiny vectors: a position plus one of four learned role vectors (foreground, background, or a box corner), with a padding marker for empty slots. The decoder reads the cached grid against those vectors and returns candidate masks with quality scores, all in a few milliseconds. Encode once, decode per click. That asymmetry is why selecting an object by clicking feels instant.

The prompt path lives in visual_processing/segmentation/segment_anything/ in the Wunjo Make repo: predictor.py for set_image and predict, modeling/prompt_encoder.py for how points and boxes become embeddings, and segment_anything/detect.py for the click-to-label wrapper. If you build your own interactive selection, the one rule that matters is: cache the embedding, then let the clicks do the cheap work.

References

Top comments (0)