MNSeg: Mamba-based 3D neuron segmentation integrated with bidirectional attention mechanism and topological loss

Segmenting neurons can effectively enhance their automatic reconstruction in 3D neuronal images, which is of great significance for revealing the structure and function of the nervous system. However, the fine nerve fibers extend over long distances in the images and the complex fiber structures of neurons are often intertwined, which pose challenges to the neuron segmentation models. In this paper, we propose a Mamba-based 3D Neuron Segmentation model, termed MNSeg, which can maintain the complete neuron topology structure and effectively distinguish different nerve fibers in segmenting the neurons. In MNSeg, we use the Mamba block to enhance the model in exploring the long-range dependencies of nerve fibers in the neuronal images. We also design a Mamba-integrated bidirectional attention (MBA) module to enhance the model's attention on the edges of the nerve fibers by assembling forward and reverse attention mechanisms with Mamba-based attention weights. Meanwhile, we apply a topological loss by combining cross-entropy loss, Dice loss, and centerline Dice loss to train MNSeg, which helps to suppress the model's excessive attention to the background area of the fiber edges in segmenting the neurons, effectively improving the model's ability to distinguish between different nerve fibers. The experimental results show that MNSeg not only improves the neuron segmentation, but also effectively maintains the neuron topological integrity, distinguishes the edges of different nerve fibers, and significantly enhances the performance of neuron reconstruction. The code is available at https://github.com/cmypromising/MNSeg.