Manifold-Aware Local Feature Modeling for Semi-Supervised Medical Image Segmentation

TL;DR

The Manifold-Aware Local Feature Modeling Network (MANet) is introduced, which enhances the U-Net architecture by incorporating manifold supervision signals and shows improved generalization across various semi-supervised segmentation networks, highlighting its robustness and effectiveness.

Abstract

Achieving precise medical image segmentation is vital for effective treatment planning and accurate disease diagnosis. Traditional fully-supervised deep learning methods, though highly precise, are heavily reliant on large volumes of labeled data, which are often difficult to obtain due to the expertise required for medical annotations. This has led to the rise of semi-supervised learning approaches that utilize both labeled and unlabeled data to mitigate the label scarcity issue. In this paper, we introduce the Manifold-Aware Local Feature Modeling Network (MANet), which enhances the U-Net architecture by incorporating manifold supervision signals. This approach focuses on improving boundary accuracy, which is crucial for reliable medical diagnosis. To further extend the versatility of our method, we propose two variants: MA-Sobel and MA-Canny. The MA-Sobel variant employs the Sobel operator, which is effective for both 2D and 3D data, while the MA-Canny variant utilizes the Canny operator, specifically designed for 2D images, to refine boundary detection. These variants allow our method to adapt to various medical image modalities and dimensionalities, ensuring broader applicability. Our extensive experiments on datasets such as ACDC, LA, and Pancreas-NIH demonstrate that MANet consistently surpasses state-of-the-art methods in performance metrics like Dice and Jaccard scores. The proposed method also shows improved generalization across various semi-supervised segmentation networks, highlighting its robustness and effectiveness. Visual analysis of segmentation results confirms that MANet offers clearer and more accurate class boundaries, underscoring the value of manifold information in medical image segmentation.

Figures (3)

• Figure 1: Visual demonstration of the importance of manifold supervision in medical imaging. (a) Input image from the ACDC dataset. (b1) Baseline prediction with segmentation supervision. (c1) The supervision signal of the baseline method, including the ground truth for labeled data and pseudo-labels for unlabeled data. (b2) Our prediction with both segmentation and manifold supervision. (c2) The manifold supervision signal. The boundary areas in predictions with manifold supervision are clearer than those in results with segmentation supervision alone.
• Figure 2: The overall network architecture. In the figure, labeled and unlabeled data are first concatenated batch-wise ($\textcircled{c}$) and then fed simultaneously into the encoder to extract feature representations. These features are passed into the decoder, which consists of two branches: the base branch and the manifold branch. The final layer of each branch is a distinct output head, with the base branch utilizing a segmentation head and the manifold branch employing a manifold head. These heads generate predictions, which are then split batch-wise ($\textcircled{s}$) into labeled and unlabeled portions, supervised by their respective segmentation and manifold signals. Specifically, for unlabeled data, the segmentation supervision is derived from the pseudo-labels generated by the pseudo-label generation network. Meanwhile, the manifold supervision is generated by the manifold generator, which uses the corresponding segmentation signals. During the testing phase, the manifold branch can be removed, allowing the network to perform inference using only the base branch (highlighted by the blue modules), thereby incurring no additional computational cost compared to the baseline method.
• Figure 3: The visualization results from the LA dataset (left) and the ACDC dataset (right), with both datasets using 10% labeled data. The baseline method employed is BCP bai2023bidirectional.