CelloType: A Unified Model for Segmentation and Classification of Tissue Images

Abstract

Cell segmentation and classification are critical tasks in spatial omics data analysis. We introduce CelloType, an end-to-end model designed for cell segmentation and classification of biomedical microscopy images. Unlike the traditional two-stage approach of segmentation followed by classification, CelloType adopts a multi-task learning approach that connects the segmentation and classification tasks and simultaneously boost the performance of both tasks. CelloType leverages Transformer-based deep learning techniques for enhanced accuracy of object detection, segmentation, and classification. It outperforms existing segmentation methods using ground-truths from public databases. In terms of classification, CelloType outperforms a baseline model comprised of state-of-the-art methods for individual tasks. Using multiplexed tissue images, we further demonstrate the utility of CelloType for multi-scale segmentation and classification of both cellular and non-cellular elements in a tissue. The enhanced accuracy and multi-task-learning ability of CelloType facilitate automated annotation of rapidly growing spatial omics data.

Figure 1 –. Overview of CelloType.

a) Overall architecture, input, and output of CelloType. First, a Transformer-based feature extractor is employed to derive multi-scale features ($C_b$) from the image. Second, using a Transformer-based architecture, the DINO object detection module extracts latent features ($C_e$) and query embeddings ($q_c$) that are combined to generate object detection boxes with cell type labels. Subsequently, the MaskDINO module integrates the extracted image features with DINO’s outputs, resulting in detailed instance segmentation and cell type classification. During training, the model is optimized based on an overall loss function ($Loss$) that considers losses based on cell segmentation mask ${(\lambda }_{mask}{L}_{mask})$, bounding box (${\lambda }_{box}{L}_{box}$), and cell type label $\left({\lambda }_{cls}{L}_{cls}\right)$. b) Input, output, and architecture of the DINO module. The DINO module consists of a multi-layer Transformer and multiple prediction heads. DINO starts by flattening the multi-scale features from the Transformer-based feature extractor. These features are merged with positional embeddings to preserve spatial context (step 1 in the figure). DINO then employs a mixed query selection strategy, initializing positional queries $\left(Q_{pos}\right)$ as anchor detection boxes and maintaining content queries $\left(Q_{content}\right)$ as learnable features, thus adapting to the diverse characteristics of cells (step 2). The model refines these anchor boxes through decoder layers using deformable attention mechanism and employs contrastive denoising training by introducing noise to ground truth (GT) labels and boxes to improve robustness and accuracy. Then a linear projection acts as the classification branch to produce the classification results for each box (step 3). c) Multi-scale ability of CelloType. CelloType is versatile and can perform a range of end-to-end tasks at different scales, including cell segmentation, nuclear segmentation, microanatomical structure segmentation, and full instance segmentation with corresponding class annotations.

Figure 2 –. Evaluation of segmentation accuracy using TissueNet datasets

a) Average Precision (AP) across Intersection over Union (IoU) thresholds for cell segmentation by Mesmer, Cellpose2, CelloType and CelloType_C (CelloType with confidence score). Mean AP value across IoU thresholds of 0.5–0.9 (mAP) for each method is indicated in parenthesis. b) AP across IoU thresholds for nuclear segmentation. c) Performance of methods stratified by imaging platform and tissue type. The top left heatmap shows the mAP scores for cell segmentation stratified by imaging platform, including CODEX, CyCIF, IMC, MIBI, MxIF and Vertra. The top right heatmap shows the mAP scores for cell segmentation stratified by tissue type, including breast, gastrointestinal, immune, pancreas and skin. The second row of heatmaps shows the mAP values for nuclear segmentation. d) Representative examples of cell segmentation of immune tissue imaged using Vectra platform. Blue, nuclear channel; green,membrane channel; white, cell boundary. The red box highlights a representative region that the methods perform differently. The AP75 score (Average precision at IoU threshold of 0.75) is displayed on the images. e) Representative examples of nuclear segmentation of gastrointestinal tissue using the IMC platform. The AP50 scores are shown on the images.

Figure 3 –. Evaluation of segmentation accuracy using Cellpose Cyto dataset

a) Average precision (AP) across Intersection over Union (IoU) thresholds for Cellpose2, CelloType and CelloType_C (CelloType with confidence score). Mean AP value across IoU thresholds of 0.5–0.9 (mAP) for each method is indicated in parenthesis. b) Mean AP values of Cellpose2, CelloType, and CelloType_C stratified by imaging modalities and cell types. The test dataset comprises microscopy and non-microscopy images from the Cellpose Cyto dataset that comprises 6 subsets, including Cells (Cell Image Library), Cells (Fluorecent), Cells (Non-fluorecent), Cells (Membrane), Other microscopy, and Non-microscopy. c) Representative examples of cell segmentation of a microscopy image by the compared methods. The red boxes highlight a representative region that the methods perform differently. The AP75 score is displayed on the images. d) Representative examples of cell segmentation of a non-fluorescent image by the compared methods.

Figure 4 –. CelloType performs joint segmentation and cell type classification.

a) Barplot showing AP50 values for cell type annotation by the two compared methods. b) Line plot showing the relationship between classification accuracy and confidence score threshold by the two methods. c) Representative examples of cell segmentation and classification results using the colorectal cancer CODEX dataset. Each row represents a 200 by 200 pixels field of view (FOV) of a CODEX image. Each FOV shows predicted cell segmentation masks (boxes) and cell types (colors). Ground Truth, manually annotated cell types; CelloType, end-to-end cell segmentation and cell type classification; Cellpose2+CellSighter, cell segmentation by Cellpose 2 followed by cell type classification by CellSighter. Randomly selected confidence scores for cell classification computed by the two methods were displayed next to the predicted instances.

Figure 5 –. Performance benchmarking of Cellpose2 and CellSighter.

Each method was evaluated for its originally intended task, namely Cellpose2 for segmentation and CellSighter for cell classification. Colorectal cancer CODEX dataset was used for benchmarking purpose. a) AP value of segmentation across a range of IoU thresholds. Mean AP value (mAP) is shown in parenthesis. b) Heatmap showing the confusion matrix of CellSighter cell type classification results. Ground truth cell segmentation masks were used as input to CellSighter. Each grid in the heatmap includes an accuracy score and the count of cells. c) Barplot showing the precision scores for each class identified by the CellSighter model based on the ground truth cell segmentation mask, with an overall mean precision of 0.53.

Figure 6 –. CelloType supports joint multi-scale segmentation and classification.

a) Performance evaluation of CelloType stratified by cell and microanatomic structure types. The bar plot shows the mean and 95% confidence interval of AP50 values in 5-fold cross-validation experiments. b) Line plot showing the relationship between classification accuracy and confidence score threshold. c) Representative examples of multi-scale segmentation and classification using human bone marrow CODEX data. The first row of images shows an example of bone marrow area consisting of various types of smaller hematopoietic cells and much larger adipocytes. The second row of images shows an example of bone marrow area consisting of various hematopoietic cell types and microanatomic structure such as trabecula bone fragments. Randomly selected confidence scores for cell classification were displayed next to the predicted instances.