Keypoint-Augmented Self-Supervised Learning for Medical Image Segmentation with Limited Annotation
Pretraining CNN models (i.e., UNet) through self-supervision has become a powerful approach to facilitate medical image segmentation under low annotation regimes. Recent contrastive learning methods encourage similar global representations when the same image undergoes different transformations, or...
Saved in:
Main Authors | , , , , |
---|---|
Format | Journal Article |
Language | English |
Published |
02.10.2023
|
Subjects | |
Online Access | Get full text |
Cover
Loading…
Summary: | Pretraining CNN models (i.e., UNet) through self-supervision has become a
powerful approach to facilitate medical image segmentation under low annotation
regimes. Recent contrastive learning methods encourage similar global
representations when the same image undergoes different transformations, or
enforce invariance across different image/patch features that are intrinsically
correlated. However, CNN-extracted global and local features are limited in
capturing long-range spatial dependencies that are essential in biological
anatomy. To this end, we present a keypoint-augmented fusion layer that
extracts representations preserving both short- and long-range self-attention.
In particular, we augment the CNN feature map at multiple scales by
incorporating an additional input that learns long-range spatial self-attention
among localized keypoint features. Further, we introduce both global and local
self-supervised pretraining for the framework. At the global scale, we obtain
global representations from both the bottleneck of the UNet, and by aggregating
multiscale keypoint features. These global features are subsequently
regularized through image-level contrastive objectives. At the local scale, we
define a distance-based criterion to first establish correspondences among
keypoints and encourage similarity between their features. Through extensive
experiments on both MRI and CT segmentation tasks, we demonstrate the
architectural advantages of our proposed method in comparison to both CNN and
Transformer-based UNets, when all architectures are trained with randomly
initialized weights. With our proposed pretraining strategy, our method further
outperforms existing SSL methods by producing more robust self-attention and
achieving state-of-the-art segmentation results. The code is available at
https://github.com/zshyang/kaf.git. |
---|---|
DOI: | 10.48550/arxiv.2310.01680 |