Promptable segmentation

Overview

A promptable segmenter takes an input and a prompt, such as a click or a box on the target, and returns the mask of what the prompt points at. Two carded models do this for medical images, and a third does it on 3D scenes. They differ in how much they are adapted:

• A 2D model, SAM, used as released. The authors evaluate it zero-shot on 19 medical datasets with point and box prompts, with no training described 2304.10517 §2.1 2304.10517 §2.2.
• A 3D model, SAM-Med3D, trained from scratch on 3D scans. Its architecture has a 3D image encoder, a point-prompt encoder and a mask decoder 2310.15161 §4.1. The shipped code runs the encoder once per volume and reruns only the prompt encoder and mask decoder for each click (card: Paper vs code and Pitfalls).
• A 3D feature field, Feature 3DGS, in which SAM's decoder reads features rendered from a trained scene, so new camera views skip SAM's encoder 2312.03203 §4.2.

Two more cards are listed under this concept, and neither is a prompted segmenter. Gaussian Grouping runs SAM in automatic mode to supervise its 3D labels. Its groups are anonymous IDs, and text queries need a separate detector 2312.00732 §3.2 2312.00732 §4.3. MaPa splits a mesh into segments by geometry and uses the segment mask as conditioning for a ControlNet-style diffusion model 2404.17569 §3.1. Neither is counted as a variant here.

Variants

Zero-shot 2D promptable model (SAM)

SAM is prompted with one point, one point per part, one box, or one box per part. Each prompt set is made for one image and one class 2304.10517 §2.2.1. Box prompts do best. One box per part, up to three parts, averages 0.6542 mean IoU across 28 tasks, and one box around the largest part averages 0.5891 2304.10517 §3.1 2304.10517 §3.2. With one point, the mean is 0.4595. That beats RITM (0.1322), SimpleClick (0.1910) and FocalClick (0.2240), and SAM is ahead of all three on 24 of 28 tasks 2304.10517 §3.2. Single-prompt IoU ranges from 0.1136 to 0.9118 across tasks 2304.10517 §3.1. Card: Segment Anything Model for Medical Image Analysis.

3D promptable model for volumes

A 3D Vision Transformer encodes the whole volume, and the user's clicks steer a 3D mask decoder. The first click is a random foreground voxel. Each later click is a random voxel in the error region, marked positive where the mask misses the target and negative where it marks extra tissue 2310.15161 §4.3. The model is trained from scratch on a labelled collection of 3D scans, not adapted from a 2D model 2310.15161 tab. 2 2310.15161 §4.2. Author-reported overall Dice is 76.27% with one point and 80.71% with ten 2310.15161 tab. 3. Card: SAM-Med3D.

Volumes wider than 128 voxels per axis are handled, according to the paper, by a sliding window that overlaps by half 2310.15161 §4.3. The released code does not include that window, so this step is paper-only card: Paper vs code.

Prompted segmentation in a 3D feature field

Feature 3DGS gives each 3D Gaussian a feature vector and trains it by distillation from a 2D foundation model, using the same splatting that renders colour 2312.03203 §3.1 2312.03203 §3.2. For SAM, the rendered feature goes to the SAM decoder alone, so novel views skip the SAM encoder 2312.03203 §4.2. The authors report this path as up to 1.7 times faster in total than rendering a view and running SAM's encoder and decoder, with nearly equivalent masks 2312.03203 fig. 4. The card records no accuracy for prompts in this mode. The paper's segmentation table uses the LSeg teacher and label queries, not SAM prompts 2312.03203 §4 2312.03203 tab. 2. Card: Feature 3DGS.

Pre-trained 3D encoder for semantic segmentation

The same encoder can serve as the pre-trained start of a conventional segmenter. In the paper's UNETR experiment, using the encoder raised the average Dice across four datasets from 83.09 to 85.79, with the largest gain of 5.63 points on AMOS 2310.15161 tab. 6 2310.15161 §5.1.6. This variant needs labelled data for the target task, because it still trains a decoder for that task. Card: SAM-Med3D.

When to use which

The one-click and per-slice comparison is the SAM-Med3D paper's own headline and pits a single 3D click against a click for every slice. Expect the gap to be smaller or larger on other data, and the baseline weights are not stated 2310.15161 §4.3.

Implementation sketch

The loop below follows the click rules and the decoder call of the 3D model 2310.15161 §4.1 2310.15161 §4.3. It is our sketch, not the paper's algorithm, which the paper does not give. The 2D SAM evaluation uses a different prompt construction, described under Variants and in 2304.10517 §2.2.1.

# Interactive segmentation of one ROI  [arxiv:2310.15161 §4.3]
# Input:  volume (D, H, W), z-scored; gt (D, H, W), used only to simulate clicks
# Output: binary mask (D, H, W)

def segment_interactive(model, volume, gt, num_clicks):
    # Encode the volume once; every click reuses these embeddings
    emb = model.image_encoder(volume)
    points = []
    prev_low_res = None
    mask = zeros_like(gt)
    for k in range(num_clicks):
        # First click: random foreground voxel. Later: random error voxel
        p, label = next_click(mask, gt, first=(k == 0))
        points.append((p, label))
        # Decoder reruns with every click so far and the previous logits
        low_res = model.mask_decoder(emb, points, prev_low_res)
        prev_low_res = low_res
        # Logit above 0 is probability above 0.5; this feeds the next click
        mask = upsample(low_res, gt.shape) > 0
    return mask

Points where the variants differ or can break:

• Use the array axis order of your volume for click coordinates. The card's Pitfalls record that the released code takes (D, H, W) indices, not NIfTI (x, y, z) order.
• The 3D encoder's position embedding is fixed to the 128-voxel ROI. Another ROI size needs the embedding resized card: Pitfalls.
• The paper does not give the click rule used during training. Keep one rule for training and evaluation (see Contradictions).
• The 2D SAM prompts are tight boxes or interior points taken from the ground truth. The paper says it simulates a user and does not measure real prompts 2304.10517 §2.2.1.
• The paper does not say how the outputs of several box prompts are merged. The released code combines them by a pixel-wise OR card: Paper vs code.
• The 2D SAM code rescales each image by its own maximum and converts it to RGB. The paper does not state this step, and it changes the contrast SAM sees on CT and PET converted to PNG card: Pitfalls.
• In a feature field, the SAM teacher runs on a padded 1024 by 1024 input, so the content of a non-square image fills only part of its 64 by 64 map. Pad the rendered map back to 64 by 64 before the SAM decoder runs 2312.03203 app. C card: Pitfalls.

Contradictions

• Click rule in training and evaluation. The paper does not say which rule its reported numbers use 2310.15161 §4.3. The released code keeps every click at evaluation, but during training it sends only the newest click unless a flag is set card: Pitfalls. Both sides are reported; neither is resolved here.
• Training length. The paper reports 800 epochs for its first training stage 2310.15161 §4.2. The training default in the release is 200, and the release has no first-stage config card: Paper vs code.
• Number of clicks evaluated. The paper reports 1, 3, 5 and 10 points 2310.15161 tab. 3. The released evaluation scripts run one click card: Paper vs code.
• Brain Dice. The 10-point model scores 61.27 on the Brain column of one table 2310.15161 tab. 5 and 61.91 in the "all" row of another 2310.15161 tab. 7. The two tables use different validation sets, so do not treat them as one series.
• Next point for iterative prompts. The SAM paper's text places the next point at the central location of the largest error component 2304.10517 §2.2.1. Its Algorithm 1 uses the point farthest from the error region's boundary, with random ties 2304.10517 alg. 1. The card follows Algorithm 1 and the code. Both sides are reported here.
• Single-prompt SAM on spine MRI. The SAM paper's abstract gives 0.1135 and its §3.1 gives 0.1136 for the same task 2304.10517 §3.1.
• Zero-shot SAM across the two papers. SAM-Med3D reports SAM at 16.15 Dice with N points per volume on its 16 validation datasets 2310.15161 tab. 3. The SAM paper reports 0.4595 mean IoU with one point on 28 tasks 2304.10517 §3.2. The data, metric and prompt counts all differ, so the two numbers are not comparable.
• Source of all numbers. Every number here is author-reported 2310.15161 §5.1 2304.10517 §3. No independent reproduction is recorded in this wiki.

Papers

• Segment Anything Model for Medical Image Analysis: an Experimental Study: zero-shot evaluation of SAM under five prompt modes on 19 medical datasets.
• SAM-Med3D: a 3D promptable model for volumetric medical scans, with a pre-trained encoder shown on semantic segmentation.
• Feature 3DGS: Supercharging 3D Gaussian Splatting to Enable Distilled Feature Fields: a feature field distilled from SAM and LSeg, with a speed-up for rendering the features.
• Gaussian Grouping: Segment and Edit Anything in 3D Scenes: per-Gaussian identity labels from SAM masks tracked across views. It is listed here for those labels, not as a prompted segmenter.
• MaPa: text-driven material painting for 3D meshes. It is listed here for its segment masks, not as a prompted segmenter.