Self-supervised Pretraining of Cell Segmentation Models

1. Problem Statement

Cell instance segmentation is critical for quantifying cell morphology in research and clinical settings (e.g., predicting cell health, senescence, and gene expression). However, the field faces two primary bottlenecks:

• Data Scarcity: High-quality, hand-labeled microscopy datasets are scarce and expensive to produce due to the time-consuming nature of annotating tightly packed or low-contrast cells.
• Domain Shift in Foundation Models: Recent approaches utilize large-scale foundation models (like the Segment Anything Model, SAM) pretrained on natural images. However, these models encode priors based on natural objectness and texture (sharp edges, distinct backgrounds) that do not align well with microscopy data, where boundaries are often sparse, blurry, and morphologically diverse. This leads to degraded performance when applying these models to microscopy images without extensive fine-tuning.

2. Methodology: DINOCell

The authors propose DINOCell, a three-stage training framework designed to bridge the gap between natural image representations and microscopy data through self-supervised domain adaptation.

A. Data Strategy

• Domain Adaptation Dataset: A curated set of 130,000 unlabeled cell images from 63 public sources (e.g., Broad Institute Cellpainting Gallery, BBBC). This dataset covers 50+ cell types across four modalities: Phase-Contrast (PC), Brightfield (BF), Differential Interference Contrast (DIC), and Fluorescence. Images are processed to focus on high-density regions.
• Fine-Tuning Dataset: 15 publicly available annotated datasets containing ~9,000 images with ground-truth instance masks.
• Data Cleaning: The authors implemented "Ignore Masking" to handle missing annotations (common in public datasets). Pixels without ground truth labels are down-weighted (weight = 0.05) during loss calculation to prevent the model from learning false negatives or overfitting to incomplete labels.

B. Architecture

• Encoder: A Vision Transformer (ViT-B) initialized with DINOv2 weights. The patch size is reduced from 14x14 to 8x8 to capture finer cellular details.
• Decoder: A lightweight upsampling decoder consisting of transposed convolutions and upsampling blocks.
• Objective: Instead of predicting masks directly, the model predicts flow fields (inspired by Cellpose). It outputs vector fields ($dx, dy$) pointing toward cell centers and a cell probability map. Instance masks are derived by integrating these flows to find fixed points.

C. Three-Stage Training Pipeline

1. Self-Supervised Pretraining: The model starts with weights from the general-purpose DINOv2 (ViT-B) model pretrained on 1.2 billion natural images.
2. Domain Adaptation (Self-Supervised): The encoder is further trained on the 130k unlabeled microscopy images using the DINOv2 self-distillation objective (student-teacher framework with iBOT). This step aligns the visual representations with cellular morphology without requiring labels.
3. Supervised Fine-Tuning: The full model (encoder + decoder) is trained on the labeled datasets to predict flow fields and cell probabilities. The encoder is unfrozen, and a consistent learning rate is used.

3. Key Contributions

• Domain-Adapted Self-Supervised Pretraining: Demonstrates that adapting a general foundation model (DINOv2) to the specific domain of microscopy via self-supervised learning is more effective than using task-specific priors from models like SAM.
• Ignore Masking Technique: A robust mechanism to handle noisy public datasets with missing annotations, significantly improving training stability and generalization.
• High-Resolution Patching: Modifying the ViT patch size from 14x14 to 8x8 allows the model to capture the fine-grained structural details necessary for cell segmentation.
• Comprehensive Evaluation: Introduces the Maximum Matching Accuracy (MMA) metric and evaluates models on both in-distribution (LIVECell) and challenging out-of-distribution (zero-shot) datasets.

4. Results

The model was evaluated against leading baselines: SAMCell (SAM-initialized) and Cellpose-SAM.

• In-Distribution Performance (LIVECell Benchmark):

  • DINOCell achieved a SEG score of 0.784, representing a 10.42% improvement over the leading SAM-based model (Cellpose-SAM).
  • It achieved an MMA score of 0.875, outperforming specialist and generalist baselines.
  • Qualitatively, DINOCell showed superior resistance to "clumping" (merging multiple cells) and "under-segmentation" (missing cells).

• Zero-Shot Performance (Out-of-Distribution):

  • DINOCell demonstrated exceptional generalization on three unseen datasets (PBL-HEK, PBL-N2A, Glioma C6).
  • Glioma C6 (High Variance): DINOCell achieved an MMA of 0.737, surpassing the best baseline (SAMCell-cyto) by 43.84%. Baseline models often failed to segment any cells in this dataset due to extreme morphological variance and lighting conditions.
  • PBL-HEK: Improved MMA by 16.84% over the baseline.

• Ablation Studies:

  • Pretraining Impact: Removing pretraining resulted in a 29.51% drop in MMA on the test set and up to 285% drop on zero-shot datasets.
  • Domain Adaptation Impact: While domain adaptation improved the clarity of encoder representations (visualized via PCA), its direct impact on final segmentation metrics was modest compared to the massive gain from the initial DINOv2 pretraining. However, it provided a slight boost in generalization to highly out-of-domain data (Glioma C6).

5. Significance

This work establishes that self-supervised representation learning is a superior initialization strategy for microscopy tasks compared to direct transfer from natural-image segmentation models.

• Robustness: By learning domain-appropriate visual priors (cell shapes, textures, and low-contrast boundaries) before supervised learning, DINOCell generalizes significantly better to unseen cell types and imaging conditions.
• Efficiency: The framework reduces the reliance on massive, perfectly labeled datasets by leveraging unlabeled data for domain alignment.
• Future Direction: It highlights that for specialized scientific imaging, "domain adaptation" of foundation models is a critical step that should precede task-specific fine-tuning to overcome the limitations of natural-image priors.