NVIDIA’s VISTA-2D model revolutionizes cellular imaging and spatial omics

According to the NVIDIA Technical Blog, NVIDIA has unveiled VISTA-2D, a baseline model designed to significantly improve cell segmentation in cellular imaging and spatial omics workflows. The model aims to improve the accuracy of downstream tasks by leveraging advanced image embedding techniques.

Feature Extraction and Clustering

The VISTA-2D model uses an image encoder to generate embeddings that can be converted into segmentation masks. These embeddings provide essential information about cell morphology, enabling accurate cell segmentation. NVIDIA’s blog post explains that these embeddings can be clustered to automatically group cells with similar morphology.

To demonstrate the capabilities of the model, NVIDIA has provided a detailed Jupyter notebook that walks users through the process of segmenting cells and extracting spatial features using VISTA-2D. The notebook also shows how to cluster these features using RAPIDS, creating an automated pipeline for classifying cell types.

Prerequisites and setup

Users interested in exploring the VISTA-2D model will need a basic understanding of Python, Jupyter, and Docker. The Docker container required for this tutorial can be started with the following command:

docker run --rm -it \
    -v /path/to/this/repo/:/workspace \
    -p 8888:8888 \
    --gpus all \
    nvcr.io/nvidia/pytorch:24.03-py3 \
    /bin/bash

Additional Python packages required for the tutorial can be installed using:

pip install -r requirements.txt

Cell segmentation using VISTA-2D

The initial step is to load the VISTA-2D model checkpoint and use it to segment cells in the image. The segmentation process generates a feature vector for each cell, which contains all the information needed to analyze cell morphology. These vectors are then used in a clustering algorithm to group cells with similar features.

cell division

The segmentation function processes cell images via VISTA-2D to generate a segmentation mask that individually labels each cell, allowing for accurate feature extraction for each cell.

img_path="example_livecell_image.tif"
patch, segmentation, pred_mask = segment_cells(img_path, model_ckpt)

Split plotting

The segmented image can be visually verified using the plot_segmentation function, which displays the original image, the segmentation results, and individual masks for each cell.

plot_segmentation(patch, segmentation, pred_mask)

Clustering features using RAPIDS

Once the feature vectors are extracted, they are clustered using RAPIDS, a GPU-accelerated machine learning library. The TruncatedSVD algorithm reduces the dimensionality of the feature vectors, making it easier to visualize clusters in 3D space.

dim_red_model = TruncatedSVD(n_components=3)
X = dim_red_model.fit_transform(cell_features)

We then cluster the reduced feature vectors using the DBSCAN algorithm. This method assigns a cluster label to each cell, which can be visualized using Plotly for interactive 3D plots.

model = DBSCAN(eps=0.003, min_samples=2)
labels = model.fit_predict(X)

conclusion

NVIDIA’s VISTA-2D model provides accurate cell segmentation and feature extraction, providing significant advances in cell imaging and spatial omics. Combined with RAPIDS for clustering, this model efficiently classifies cell types, paving the way for more detailed and automated biological studies.

Image source: Shutterstock