trAIce3D: A Prompt-Driven Transformer Based U-Net for Semantic Segmentation of Microglial Cells from Large-Scale 3D Microscopy Images

2025-07-30

AI Summary

Current methods for semantic segmentation of microglia—including somata and fine processes—in large-scale 3D microscopic images suffer from low accuracy, difficulty in disentangling overlapping structures, poor noise robustness, and weak cross-dataset generalization. To address these challenges, we propose a two-stage prompt-driven framework: first, leveraging precisely localized soma coordinates as geometric prompts; second, guiding process segmentation via a vision Transformer integrated with cross-attention mechanisms. Our architecture is built upon a 3D U-Net enhanced with sliding-window inference, self-supervised pretraining, prompt learning, and cross-attention–augmented skip connections. Evaluated on a dataset comprising 41,230 cells, our method achieves significant improvements in segmentation accuracy and generalizability. It is the first to enable joint high-precision segmentation of both somata and processes. Moreover, the framework demonstrates strong extensibility, readily adapting to other neural cell types—including neurons and astrocytes—without architectural modification.

Technology Category

• Computer Vision: Segmentation
• Natural Language Processing: Prompt Engineering / Prompting
• Machine Learning: Neuro-Symbolic Learning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The shape of a cell contains essential information about its function within the biological system. Segmenting these structures from large-scale 3D microscopy images is challenging, limiting clinical insights especially for microglia, immune-associated cells involved in neurodegenerative diseases. Existing segmentation methods mainly focus on cell bodies, struggle with overlapping structures, perform poorly on noisy images, require hyperparameter tuning for each new dataset, or rely on tedious semi-automated approaches. We introduce trAIce3D, a deep-learning architecture designed for precise microglia segmentation, capturing both somas and branches. It employs a two-stage approach: first, a 3D U-Net with vision transformers in the encoder detects somas using a sliding-window technique to cover the entire image. Then, the same architecture, enhanced with cross-attention blocks in skip connections, refines each soma and its branches by using soma coordinates as a prompt and a 3D window around the target cell as input. Training occurs in two phases: self-supervised Soma Segmentation, followed by prompt-based Branch Segmentation, leveraging pre-trained weights from the first phase. Trained and evaluated on a dataset of 41,230 microglial cells, trAIce3D significantly improves segmentation accuracy and generalization, enabling scalable analysis of complex cellular morphologies. While optimized for microglia, its architecture can extend to other intricate cell types, such as neurons and astrocytes, broadening its impact on neurobiological research.

Problem

• Segment microglial cells from 3D microscopy images
• Overcome challenges in overlapping structures and noise
• Eliminate need for dataset-specific hyperparameter tuning

Innovation

• Prompt-driven transformer enhances 3D U-Net
• Two-stage approach for soma and branch segmentation
• Self-supervised and prompt-based training phases