CellTranspose: Few-shot Domain Adaptation for Cellular Instance Segmentation

Keaton, Matthew; Zaveri, Ram; Doretto, Gianfranco

CellTranspose: Few-shot Domain Adaptation for Cellular Instance Segmentation

Matthew R Keaton, Ram J Zaveri, Gianfranco Doretto

West Virginia University
WACV 2023

Paper Code

Robust segmentation across various imaging modalities and cell types.

Abstract

Automated cellular instance segmentation is a process utilized for accelerating biological research for the past two decades, and recent advancements have produced higher quality results with less effort from the biologist. Most current endeavors focus on completely cutting the researcher out of the picture by generating highly generalized models. However, these models invariably fail when faced with novel data, distributed differently than the ones used for training. In this work, rather than approaching the problem with methods that presume the availability of large amounts of target data and computing power for retraining, we address the even greater challenge of designing an approach that requires minimal amounts of new annotated data as well as training time. We do so by designing specialized contrastive losses that leverage the few annotated samples very efficiently. A large set of results show that 3 to 5 annotations lead to models with accuracy that: 1) significantly mitigate the covariate shift effects; 2) matches or surpasses other adaptation methods; 3) even approaches the performance of methods that have been fully retrained on the target distribution. The adaptation training is only a few minutes, paving a path towards a balance between model performance, computing requirements and expert-level annotation needs.

Qualitative results

CellTranspose is able to reliably adapt to unseen data with as few as 3 labeled cells, and outperforms the generalist baseline regardless of different features of the data. In this case, z corresponds to the imaging plane, where z=16 is in focus.

CellTranspose prediction and adaptation overview

Prediction process. Similar to Cellpose, our approach generates a per-pixel cell prediction mask as well as x-y gradient flows pointing to the direction to the center of the cell, which enables individual cell instances to be predicted.

Adaptation process. Using source-target domain sample pairs, predictions are updated utilizing specialized contrastive losses for both the mask and flow predictions.

Mask-based Contrastive Loss. On a per-pixel basis, the model prediction is updated to ensure predictions from the target samples are representationally aligned to source samples when belonging to the same class (cell/non-cell), and otherwise distant from each other.

2 examples of our Contrastive Flow Loss, with relative errors between the prediction (blue) and both positive (yellow) and negative (gray) source dataset samples represented by line weight. For each pixel corresponding to a cell in the target sample, we identify a location from the source sample which possesses a gradient flow vector with the highest cosine similarity, and utilize a hard-mining strategy to select multiple negative samples which are roughly a fixed cosine distance away from the target label. Similarly to the Mask-based Contrastive Loss, we utilize a loss function that aligns the target sample representation with that of the positive source sample, while moving away from each negative sample.

Comparative results with generalist approaches

CellTranspose is able to adapt to a number of imaging and tissue types, in some cases significantly outperforming the "generalist" Cellpose baseline as well as another generalist approach published alongside the TissueNet dataset. With only 3 target samples, CellTranspose in most cases even approaches the upper bound defined by training the original Cellpose model directly on each individual subtype.

Comparative results with unsupervised method

We also compare our approach to unsupervised methods which aim to utilize the entire target dataset without any labels. A clear comparison is difficult, as factors including the size of the target dataset and amount of data variability have different effects on either approach. In this case, more training samples are required to perform as well as the unsupervised method, which may in part be due to the fact that the target training data and test data come from different tissue types.

BibTeX

@inproceedings{keaton2023celltranspose,
        title={Celltranspose: Few-shot domain adaptation for cellular instance segmentation},
        author={Keaton, Matthew R and Zaveri, Ram J and Doretto, Gianfranco},
        booktitle={Proceedings of the IEEE/CVF winter conference on applications of computer vision},
        pages={455-466},
        year={2023},
        url={https://openaccess.thecvf.com/content/WACV2023/html/Keaton_CellTranspose_Few-Shot_Domain_Adaptation_for_Cellular_Instance_Segmentation_WACV_2023_paper.html}
    }