Category Prompt Mamba Network for Nuclei Segmentation and Classification

TL;DR

This work tackles patch-induced edge misalignment and severe class imbalance in nuclei segmentation and classification by proposing CP-Mamba, a Mamba-based encoder–decoder that trains directly on full-scale images. It introduces category prompts for per-class feature learning and a probability-guided sequence sorting strategy to prioritize high-confidence information, improving representation for rare nuclei categories. The approach delivers state-of-the-art performance on four public datasets with reduced inference time, demonstrating the practicality of full-image training for large pathology images. The method promises stronger tumor immune microenvironment analysis by boosting accuracy for underrepresented nuclear categories and enabling efficient, scalable analysis in clinical pipelines.

Abstract

Nuclei segmentation and classification provide an essential basis for tumor immune microenvironment analysis. The previous nuclei segmentation and classification models require splitting large images into smaller patches for training, leading to two significant issues. First, nuclei at the borders of adjacent patches often misalign during inference. Second, this patch-based approach significantly increases the model's training and inference time. Recently, Mamba has garnered attention for its ability to model large-scale images with linear time complexity and low memory consumption. It offers a promising solution for training nuclei segmentation and classification models on full-sized images. However, the Mamba orientation-based scanning method lacks account for category-specific features, resulting in sub-optimal performance in scenarios with imbalanced class distributions. To address these challenges, this paper introduces a novel scanning strategy based on category probability sorting, which independently ranks and scans features for each category according to confidence from high to low. This approach enhances the feature representation of uncertain samples and mitigates the issues caused by imbalanced distributions. Extensive experiments conducted on four public datasets demonstrate that our method outperforms state-of-the-art approaches, delivering superior performance in nuclei segmentation and classification tasks.

Figures (6)

• Figure 1: The existing nuclei analysis framework. The images with black box represent current training samples, and the images with red box do not participate in training. (a) represents convolution network independently trains each image; (b) represents Transformer structure has quadratic computational complexity; (c) represents state space model only considers the preceding samples; (d) is our proposed category prompt network, which utilizes the classification probability as a basis to guide the sequences sorting.
• Figure 2: The framework of our proposed CP-Mamba. The method employs the category supervision information to learn each type nuclear prototype. Meantime, the probability predictions provide the guides for feature sequence ordering.
• Figure 3: The category sorting method guided by probability prediction. "i" represents the sorting by index.
• Figure 4: The model complexity analysis. The "green" bubbles represent the segmentation model, and the "red" bubbles represent the classification model. The size of the bubble represents the size of the model parameters.
• Figure 5: The visualization comparison on the segmentation task. The black boxes highlight the segmentation differences.