PA-Seg: Learning from Point Annotations for 3D Medical Image Segmentation using Contextual Regularization and Cross Knowledge Distillation

TL;DR

This paper introduces PA-Seg, a weakly supervised 3D medical image segmentation framework that uses minimal point annotations, contextual regularization, and cross knowledge distillation to achieve near fully supervised performance.

Contribution

The paper presents a novel two-stage weakly supervised learning framework for 3D medical image segmentation using only seven point annotations and introduces new regularization and distillation strategies.

Findings

Outperforms existing weakly supervised methods on public datasets.

Achieves performance close to fully supervised models after additional training.

Effective in segmenting Vestibular Schwannoma and Brain Tumors with minimal annotations.

Abstract

The success of Convolutional Neural Networks (CNNs) in 3D medical image segmentation relies on massive fully annotated 3D volumes for training that are time-consuming and labor-intensive to acquire. In this paper, we propose to annotate a segmentation target with only seven points in 3D medical images, and design a two-stage weakly supervised learning framework PA-Seg. In the first stage, we employ geodesic distance transform to expand the seed points to provide more supervision signal. To further deal with unannotated image regions during training, we propose two contextual regularization strategies, i.e., multi-view Conditional Random Field (mCRF) loss and Variance Minimization (VM) loss, where the first one encourages pixels with similar features to have consistent labels, and the second one minimizes the intensity variance for the segmented foreground and background, respectively. In the second stage, we use predictions obtained by the model pre-trained in the first stage as pseudo labels. To overcome noises in the pseudo labels, we introduce a Self and Cross Monitoring (SCM) strategy, which combines self-training with Cross Knowledge Distillation (CKD) between a primary model and an auxiliary model that learn from soft labels generated by each other. Experiments on public datasets for Vestibular Schwannoma (VS) segmentation and Brain Tumor Segmentation (BraTS) demonstrated that our model trained in the first stage outperformed existing state-of-the-art weakly supervised approaches by a large margin, and after using SCM for additional training, the model's performance was close to its fully supervised counterpart on the BraTS dataset.