Cellpose 2.0: how to train your own model

Abstract

Pretrained neural network models for biological segmentation can provide good out-of-the-box results for many image types. However, such models do not allow users to adapt the segmentation style to their specific needs and can perform suboptimally for test images that are very different from the training images. Here we introduce Cellpose 2.0, a new package that includes an ensemble of diverse pretrained models as well as a human-in-the-loop pipeline for rapid prototyping of new custom models. We show that models pretrained on the Cellpose dataset can be fine-tuned with only 500-1,000 user-annotated regions of interest (ROI) to perform nearly as well as models trained on entire datasets with up to 200,000 ROI. A human-in-the-loop approach further reduced the required user annotation to 100-200 ROI, while maintaining high-quality segmentations. We provide software tools such as an annotation graphical user interface, a model zoo and a human-in-the-loop pipeline to facilitate the adoption of Cellpose 2.0.

Fig. 1. Diverse annotation styles across ground-truth datasets.

These are examples of images that the human annotators chose to segment a certain way, where another equally valid segmentation style exists. All these examples were chosen to be representative of large categories of images in their respective datasets. a, Annotation examples from the Cellpose dataset. From left to right, these show: (i) nuclei without cytoplasm are not labeled, (ii) diffuse processes are not labeled, (iii) outlines biased toward the outside of cells and (iv) dense areas with unclear boundaries are nonetheless segmented. b, Annotation examples from the TissueNet dataset. These illustrate: (i) outlines follow membrane/cytoplasm for some image types, and include nuclei without a green channel label, (ii) outlines do not follow cytoplasm for other image types, (iii) slightly out of focus cells are not segmented and (iv) outlines drawn just around nucleus for some image types. c, Annotation examples from the LiveCell dataset. These illustrate: (i) dense labeling for some image types, (ii) no labeling in dense areas for other image types, (iii) same as (ii) and (iv) no labeling in some image areas for unknown reasons.

Fig. 2. An ensemble of models with different segmentation styles.

a, t-SNE display of the segmentation styles of images from the Cellpose, LiveCell and TissueNet datasets. The style vector computed by the neural network was embedded in two-dimensions using t-SNE and clustered into nine groups using the Leiden algorithm. Each color indicated one cluster, with the name chosen based on the most popular image category in the cluster. b, Example images from each of the nine clusters corresponding to different segmentation styles. c, Improvement of the generalist ensemble model compared to a single generalist model. d, Examples of six different images from the test set segmented with two different styles each. Error bars represent the s.e.m. across test images.

Fig. 3. State-of-the-art cellular segmentation does not require big data.

a, Segmentation of the same test image by models trained with incrementally more images and initialized from the pretrained Cellpose 1.0 model. The image category is breast, vectra from the TissueNet dataset. b, Average precision of the models as a function of the number of training masks. Shown is the performance of models initialized from the Cellpose parameters or initialized from scratch. We also show the performance of the Mesmer model, which was trained on the entire TissueNet dataset. c,d, Same as a,b for image category A172 from the LiveCell dataset. The LiveCell model is shown as a baseline, with the caveat that this model was trained to report overlapping ROI (Methods). e, Left shows the average precision curves for all image categories in the TissueNet dataset. Right shows a zoom-in for less than 3,000 training masks. f, Same as e for the LiveCell image categories.

Fig. 4. A human-in-the-loop approach for training specialized Cellpose models.

a, Schematic of human-in-the-loop procedure. This workflow is available in the Cellpose 2.0 GUI. b, A new TissueNet model on the breast, vectra category was built by sequentially annotating the five training images shown. After each image, the Cellpose model was retrained using all images annotated so far and initialized with the Cellpose parameters. On each new image, the latest model was applied and the human curator only added the ROI that were missed or incorrectly segmented by the automated method. The yellow outlines correspond to cells correctly identified by the model, and the purple outlines correspond to the new cells added by the human annotator. c, Same as b for training a LiveCell model on the A172 category.

Fig. 5. Human-in-the-loop models require minimal human annotation.

a, Test image segmentations of four models trained on the five TissueNet images from Fig. 4b with different annotation strategies. Annotations were either produced with a human-in-the-loop approach (online), or by independently annotating each image without automated help (offline). The models were either pretrained (cellpose_init) or initialized from scratch. Purple outlines correspond to the ground-truth provided by the same annotator. Yellow outlines correspond to model predictions. b, Within-human agreement was measured by having the human annotator segment the same test images twice. For the second annotation, the images were mirrored horizontally and vertically to reduce memory effects. c, Total number of manually segmented ROI for each annotation strategy. d, Average precision at an IoU of 0.5 as a function of the number of training images. e, Average precision curves as a function of the number of manually annotated ROI. f–j, Same as a–e for the image category A172 from the LiveCell dataset. All models were trained on the images from Fig. 4b, with the same annotation strategies.

Extended Data Fig. 1. Specialization of Cellpose model without trained style.

a, t-SNE embedding of segmentation styles for each image, colored according to cluster identity. b, Representative example images from each class. c, AP improvement of the model ensemble over a single generalist model. Error bars represent the standard error of the mean across test images.

Extended Data Fig. 2. Specialization of the pretrained model for images of nuclei.

a, t-SNE embedding of segmentation styles for each image, colored according to cluster identity. b, Representative example images from each class. c, AP improvement of the model ensemble over a single generalist model. Error bars represent the standard error of the mean across test images.

Extended Data Fig. 3. Segmentation performance for different metrics.

a, Average precision of segmentation on the TissueNet test set for additional IoU thresholds (IoU threshold = 0.5 also shown in Fig. 3b). Performance shown as a function of the number of training ROIs for models initialized from the pretrained Cellpose model. (left) Average precision as a function of training ROIs at three IoU thresholds; (right) average precision as a function of IoU threshold for different numbers of training ROIs. b, False negative rates on the TissueNet test set. (left) False negative rates as a function of training ROIs at three IoU thresholds; (right) false negative rates as a function of IoU thresholds for different numbers of training ROIs. c, Same as (b) for the false positives rates. d-f Same as (a-c) for the LiveCell dataset.

Extended Data Fig. 4. Models pretrained on Cellpose dataset outperform models pretrained on other datasets.

a, Average precision at IoU threshold 0.5 on the TissueNet test set for the Cellpose model pretrained on the LiveCell dataset versus pretrained on the Cellpose dataset. b, Average precision on the TissueNet test set as a function of the number of training ROIs, for models 1) trained from scratch, 2) pretrained with the Cellpose dataset (same as Fig. 3b), or 3) pretrained with the LiveCell dataset. c, Average precision on the LiveCell test set for the Cellpose model pretrained on the TissueNet dataset versus pretrained on the Cellpose dataset. d, Average precision at IoU threshold on the LiveCell test set as a function of the number of training ROIs from the LiveCell dataset, for models 1) trained from scratch, 2) pretrained with the Cellpose dataset (same as Fig. 3d), or 3) pretrained with the TissueNet dataset.

Extended Data Fig. 5. Test set performance as a function of learning rate and training epochs.

a, Average precision for models trained on the Tissuenet dataset, using 5 training images. b, Same as (a) for the Livecell dataset.