PaperlayerView exact revision on arXiv

Source-linked AI summary

VISTA3D: A Unified Segmentation Foundation Model For 3D Medical Imaging

Yufan He, Pengfei Guo, Yucheng Tang, Andriy Myronenko, Vishwesh Nath, Ziyue Xu, Dong Yang, Can Zhao, Benjamin Simon, Mason Belue, Stephanie Harmon, Baris Turkbey, Daguang Xu, Wenqi Li

arXiv:2406.05285v3cs.CV

TL;DR

3D medical segmentation needs a unified approach because task-specific models, 2D adaptations, and limited zero-shot or correction support leave important workflow gaps. VISTA3D combines 3D automatic and interactive segmentation with supervoxels distilled from 2D pretrained features. It reports strong performance across these tasks, including state-of-the-art interactive results and a 50% zero-shot performance boost.

  • Problem

    Task-specific 3D models lack zero-shot handling of novel classes and interactive correction, while 2D adaptations require substantial annotation effort and lag on complex 3D structures.

    Source: §1 · pp. 1–2

  • Method

    VISTA3D unifies 3D automatic and interactive segmentation, using 3D supervoxels generated from 2D SAM feature maps to distill image understanding.

    Source: §3.2 · p. 4, §3.3 · p. 5

  • Results

    VISTA3D achieves segmentation comparable with specialized expert models and state-of-the-art interactive results for zero-shot segmentation and editing, with zero-shot performance boosted by 50%.

    Source: §1 · p. 3, §7 · p. 8

  • Takeaways & Limitations

    The authors conclude that VISTA3D's components support an out-of-the-box human-in-the-loop workflow intended to save human effort.

    Source: §7 · p. 8

Abstract

from arXiv · show

Foundation models for interactive segmentation in 2D natural images and videos have sparked significant interest in building 3D foundation models for medical imaging. However, the domain gaps and clinical use cases for 3D medical imaging require a dedicated model that diverges from existing 2D solutions. Specifically, such foundation models should support a full workflow that can actually reduce human effort. Treating 3D medical images as sequences of 2D slices and reusing interactive 2D foundation models seems straightforward, but 2D annotation is too time-consuming for 3D tasks. Moreover, for large cohort analysis, it's the highly accurate automatic segmentation models that reduce the most human effort. However, these models lack support for interactive corrections and lack zero-shot ability for novel structures, which is a key feature of "foundation". While reusing pre-trained 2D backbones in 3D enhances zero-shot potential, their performance on complex 3D structures still lags behind leading 3D models. To address these issues, we present VISTA3D, Versatile Imaging SegmenTation and Annotation model, that targets to solve all these challenges and requirements with one unified foundation model. VISTA3D is built on top of the well-established 3D segmentation pipeline, and it is the first model to achieve state-of-the-art performance in both 3D automatic (supporting 127 classes) and 3D interactive segmentation, even when compared with top 3D expert models on large and diverse benchmarks. Additionally, VISTA3D's 3D interactive design allows efficient human correction, and a novel 3D supervoxel method that distills 2D pretrained backbones grants VISTA3D top 3D zero-shot performance. We believe the model, recipe, and insights represent a promising step towards a clinically useful 3D foundation model. Code and weights are publicly available at https://github.com/Project-MONAI/VISTA.

1. Introduction

3D medical segmentation needs accurate automation for common structures, interactive correction, and zero-shot handling of novel classes; existing 2D adaptations and task-specific models do not provide that full workflow. VISTA3D unifies these capabilities in a 3D model, with a supervoxel method and training recipe designed to support them.

1. Introduction

  • Task-specific expert models require curated training data and engineering expertise, while a model that handles varied tasks out of the box is more desirable.Source: §1 · p. 1
  • Standard anatomy models can be trained broadly, but rare pathologies and animal data remain unsupported, and clinicians also need human correction for tasks such as surgical planning.Source: §1 · p. 1
  • Applying 2D natural-image or video approaches to 3D medical scans falls short: per-slice prompts demand substantial effort, and these methods lag on complex 3D structures.The paper also notes that video tracking does not ensure spatially consistent treatment of volumetric inputs.Source: §1 · p. 2
  • A 3D medical foundation model should combine accurate automation, human refinement, zero-shot interaction, 3D context, and few-shot transfer to new classes.Source: §1 · p. 2
  • VISTA3D builds on established 3D pipelines while using insights and checkpoints from 2D models, addressing the challenge of retaining strong automatic performance across diverse classes.Source: §1 · p. 2
  • VISTA3D supports automatic and interactive segmentation across 14 challenging datasets and 127 classes, with its supervoxel method boosting zero-shot performance by 50%.The model combines a curated 11,454-scan CT dataset with a four-stage training recipe.Source: §1 · p. 3

2. Related Work

Prior work leaves gaps in zero-shot and interactive 3D segmentation; VISTA3D addresses them with supported-class and interactive branches, SAM-derived supervoxels, and staged training.

2. Related Work

  • Existing 3D foundation models cover many structures but lack zero-shot and interactive ability, while open-vocabulary approaches lag expert models and 2D prompt transfer is effort-intensive.Source: §2 · p. 3

3.1. Overview

  • VISTA3D separates 127 supported classes, segmented through a prompted automatic head, from zero-shot classes handled by 3D user clicks.The training collection contains 11,454 CT scans.Source: §3.1 · pp. 3–4

3.2. Model architecture

  • A shared SegResNet encoder feeds separate decoders; class-index embeddings prompt the automatic head, while a 3D click-attention head supports supported and novel classes.The interactive branch upsamples features for detail, then downsamples them to reduce memory and segments a local patch around each click.Source: §3.2 · p. 4
  • Interactive edits add or remove connected components containing clicks, restricting corrections to selected regions of the automatic mask.Source: §3.3 · p. 5

3.3. Data

  • The dataset combines 11,454 CT scans, manual or partial labels, 117-class TotalSegmentator pseudo-labels, and SAM-derived 3D supervoxels from slice features in three views.The supervoxels are used to train the interactive branch and provide VISTA3D zero-shot capabilities.Source: §3.3 · p. 5

3.4. Recipe

  • Four stages train and finetune the interactive branch, then train a frozen-encoder automatic branch and finetune it using synthetic data with rare anomalies.Interactive training uses manual or pseudo labels and supervoxels; later finetuning oversamples under-represented classes.Source: §3.4 · pp. 5–6

4. Supported classes results

VISTA3D’s automatic segmentation is comparable to expert baselines, and a single click can improve results. The examples show correction of false negatives and false positives, plus robust segmentation on a monkey scan.

4. Supported classes results

  • VISTA3D achieved comparable automatic performance to established baselines and much better performance when minimum human input was available.A single click was sampled from the larger false-positive or false-negative region in each sliding-window patch.Source: §4 · p. 6
  • A positive point corrects a false-negative region in the automatic liver segmentation, while a negative point removes a false-positive region.The figure shows these corrections on separate slices.Source: Figure 3 · p. 6
  • VISTA3D achieved more robust segmentation in the example monkey CT scan.The example contains two sagittal slices.Source: Figure 4 · p. 6

5. Zero-shot results

On four external datasets with novel classes, VISTA3D outperformed comparison methods in accuracy while requiring less annotation effort. Its iterative clicks refine zero-shot segmentation and affect responses in 3D space.

5. Zero-shot results

  • Figure 5 plots average dataset Dice score against the number of click points.The supplied passage identifies the axes but does not provide plotted values.Source: Figure 5 · p. 7
  • VISTA3D showed superior accuracy with reduced annotation effort on four external datasets containing classes unsupported by its automatic segmentation.The evaluation included mouse organs, kidney tumors, adrenocortical carcinoma, and hepatocellular carcinoma; the authors also report the importance of supervoxels for zero-shot ability.Source: §5 · p. 8
  • On kidney tumor, iterative clicks refine the zero-shot segmentation, with a negative point correcting over-segmentation after positive clicks.Responses on adjacent slices show that clicks affect 3D space.Source: Figure 6 · p. 7, §5 · p. 8, Figure 7 · p. 8

6. Finetuning results

Fine-tuning VISTA3D’s pretrained checkpoint outperformed baselines in few-shot settings and achieved the highest WORD Dice score among ten baselines using all 100 training cases.

6. Finetuning results

  • VISTA3D achieved the highest WORD Dice score, 0.875, among ten baselines using the full training split of 100 cases.The held-out test results also showed much better performance than the baselines under few-shot settings.Source: §6 · p. 8

7. Conclusion

VISTA3D is presented as a unified 3D CT segmentation foundation model designed for a human-in-the-loop workflow. The authors report accurate automatic segmentation, state-of-the-art interactive segmentation, and strong transfer learning, with broader modalities and clinical validation left for future work.

7. Conclusion

  • VISTA3D combines accurate automatic segmentation, state-of-the-art interactive segmentation, and strong transfer learning in a unified 3D CT foundation model.The authors say its large-scale training data, architecture, and training recipes were vital to its capabilities.Source: §7 · p. 8
  • The authors identify broader class and modality support, improved zero-shot use of natural-image datasets and checkpoints, and clinical-partner validation as future work.Planned modalities include MRI and PET.Source: §7 · p. 9

——Supplementary Material——

The supplement includes a figure showing how annotated voxels are distributed across task classes.

——Supplementary Material——

  • Figure 1 plots annotated voxels per class against class index to show the training-set distribution.Source: Figure 1 · p. 12

1. Dataset Details

The training datasets use a shared global class index to reconcile dataset-specific labels. Their spatial resolutions vary across datasets.

1. Dataset Details

  • A global index of 127 classes maps dataset-specific label indices, with a separate class-index list for each dataset.For example, the pancreas is index 1 in MSD07 and index 10 in TotalsegmentatorV2.Source: §1 · p. 12
  • Dataset spatial resolutions range from 0.45 × 0.45 × 0.45 to 1.50 × 1.50 × 7.50 mm3, with a median of 0.88 × 0.88 × 1.50 mm3.Source: §1 · p. 12
  • Table 1 summarizes the datasets used to train the model.Source: Table 1 · p. 12

2. Computational Details

VISTA3D's computational requirements vary with prompt count and image size, while its sliding-window inference is designed to stabilize memory use. The supplement reports training resource requirements and inference runtimes on a lower-end machine.

2. Computational Details

  • Training used 64 32GB NVIDIA V100 GPUs for around 20,000 total GPU hours.Source: §2 · p. 12
  • Each training iteration used 36 object-class prompts for the automatic branch and 4 for the point branch.Source: §2 · p. 12
  • Training can use 16GB-or-lower GPUs by reducing prompts per iteration, at the cost of longer training when class counts are large.Source: §2 · p. 12
  • 1m43s for VISTA3D versus 2m41s for TotalSegmentator on a typical CT scan, with 3.2s for single-click inference on the same lower-end machine.The machine had a 12GB GPU and 32GB CPU; the typical scan was 308x260x453 after resampling.Source: §2 · p. 12, Figure 2 · p. 12, Figure 3 · p. 13
  • Inference uses 128x128x128 sliding windows and adaptive CPU/GPU memory control to reduce dependence on image size and avoid out-of-memory errors.GPU memory requirements also depend on prompt number and image size.Source: §2 · p. 12

3. Additional training details

SAM2 fine-tuning converged but performed poorly on complicated 3D shapes, while VISTA3D’s training details describe point-based interactive and promptable automatic branches. The interactive sampler creates supported-class editing and zero-shot training pairs, and automatic training also uses background prompts to suppress false positives.

3.1. Stage1-Interactive branch training

  • The interactive point sampler uses points from labels or supervoxels, and adds or subtracts qualifying supervoxels to train supported-class editing and zero-shot responses.When an addition or subtraction exceeds a size limit, its training pair is also treated as zero-shot; training uses maxiter = 5.Source: §3.1 · p. 13

3.2. Stage3-Automatic branch training

  • Automatic-branch training samples class prompts from manual or pseudo-labels and includes background prompts trained to produce zero output, mitigating false positives.The training iteration uses a 128-cubic image patch, up to 32 class prompts and up to four background prompts.Source: §3.2 · p. 13

3.3. SAM2 finetuning experiments

  • SAM2 converged after 500 epochs on 8 80GB A100 GPUs but still failed to track complicated 3D shapes, despite tracking simple shapes with clear boundaries.The reported Single click performance on the MSD Task07 test set is identified in Table 2; the supplied passage gives no values.Source: §3.3 · p. 14, Figure 5 · p. 14, Table 2 · p. 14

4. Additional Results

Additional qualitative results show fine-grained editing, comparisons on hard classes, and interactive mouse-lung segmentation with responses extending beyond the clicked slice. Detailed classwise Dice scores are also reported for the test datasets.

4.1. Qualitative Results

  • VISTA3D supports pixel-level editing, including one-pixel-wide boundaries, while a bounding-box prompt cannot perform editing.The size of the edited area from one click depends on the model’s boundary understanding and may be much larger than one pixel.Source: §4.1 · p. 14
  • On held-out abdominal scans without pancreas or hepatic-vessel annotations, VISTA3D shows better detail for these hard classes than comparison examples.The scans come from the MSD task09 spleen test split, avoiding ground-truth leakage for these structures.Source: §4.1 · p. 14
  • For mouse left-lung interaction, VISTA3D responds on slices far from the click, while MedSAM lacks fine-detail editing and SegVol has low response resolution.The cited examples are from micro-CT; the far-slice response is presented as a potential way to reduce annotation effort.Source: §4.1 · p. 14

4.2. Quantitative Results

  • Detailed Dice scores are reported for all classes in the test datasets.The passage identifies Table 3 as the location of these results but provides no score values.Source: §4.2 · p. 14

5. Additional Discussions

The discussion explains why VISTA3D separates automatic and interactive branches while sharing an encoder, and gives qualitative examples of editing, cross-slice point response, and zero-shot behavior. These examples include cases where supported-class prompting can resist edits and mouse anatomy differs substantially from human anatomy.

5. Additional Discussions

  • Separating the branches avoids slower joint training and conflicts between zero-shot and automatic segmentation, while a shared encoder reduces editing inference cost and may generalize automatic features.The interactive branch is more memory intensive and can use broader training data; pilot joint training produced worse results and the automatic branch did not reach state-of-the-art performance.Source: §5 · pp. 14–15
  • One click can edit a one-pixel-wide boundary, but the resulting edited area may be much larger and depends on boundary understanding.The figure contrasts this fine-grained editing with bounding-box prompts, which cannot perform editing.Source: Figure 6 · p. 15
  • A click on slice 304 produces good responses on slices 280, 296, 312, and 328, illustrating point influence beyond the clicked slice.The example is described as potential to reduce annotation effort in 3D space.Source: Figure 10 · p. 16
  • The mouse micro-CT example highlights a scope challenge: the large structural difference from humans can defeat automatic models trained on human anatomy.The figure discusses mouse left-lung zero-shot performance.Source: Figure 11 · p. 17
  • Zero-shot embeddings help VISTA3D follow clicks for fluid outside a kidney and liver substructure edits, where SegVol or supported-class prompting can ignore them.The examples link click resistance to overfitting to common organ shapes, intensities, or locations.Source: Figure 12 · p. 17, Figure 13 · p. 17

Related papers

Loading 2406.05285v3…