[Paper Review] Deep feature transfer between localization and segmentation tasks

Szu-Yeu Hu, Andrew Beers+9 more|arXiv (Cornell University)|Nov 6, 2018

Advanced Neural Network Applications8 references4 citations

TL;DR

This paper proposes a two-stage pre-training method for U-net-based optic disc segmentation, first training the encoder as a localization network to predict the disc's center, then fine-tuning it for segmentation. The approach transfers semantic features learned during localization to improve segmentation performance with significantly fewer annotated segmentation masks, achieving a mean Dice coefficient of 0.88 with reduced annotation burden.

ABSTRACT

In this paper, we propose a new pre-training scheme for U-net based image segmentation. We first train the encoding arm as a localization network to predict the center of the target, before extending it into a U-net architecture for segmentation. We apply our proposed method to the problem of segmenting the optic disc from fundus photographs. Our work shows that the features learned by encoding arm can be transferred to the segmentation network to reduce the annotation burden. We propose that an approach could have broad utility for medical image segmentation, and alleviate the burden of delineating complex structures by pre-training on annotations that are much easier to acquire.

Motivation & Objective

• To reduce the annotation burden in medical image segmentation by leveraging easier-to-annotate localization data.
• To investigate whether features learned from a localization task can improve segmentation performance in a U-net architecture.
• To evaluate the effectiveness of pre-training on localization before fine-tuning for segmentation in retinal fundus images.
• To compare the performance of pre-trained models against randomly initialized U-nets under varying data availability.
• To explore the scalability of this approach to other medical imaging tasks with complex structures.

Proposed method

• The encoding arm of a U-net is first trained as a regression network to predict the (x,y) coordinates of the optic disc center using mean squared error loss.
• After convergence, the encoder weights are frozen and the decoder is added to form a complete U-net for segmentation.
• The segmentation network is trained using a negative log soft Dice loss function with Adam optimization on binary mask annotations.
• Skip connections are added between encoder and decoder features to preserve spatial details during up-sampling.
• The model uses 3×3 convolutions, batch normalization, ReLU activations, max-pooling, and dropout to prevent overfitting.
• A two-phase training scheme is applied: first localization, then segmentation with shared encoder features.

Experimental results

Research questions

• RQ1 Can features learned from a simple localization task improve segmentation performance in medical images?
• RQ2 Does pre-training on localization data reduce the number of required segmentation annotations?
• RQ3 How does the performance of a localization-pretrained U-net compare to a randomly initialized U-net in low-data regimes?
• RQ4 To what extent are semantic features such as shape, texture, and boundary characteristics transferable from localization to segmentation?
• RQ5 Can this approach be generalized to other medical image segmentation tasks with complex anatomical structures?

Key findings

• The proposed method achieved a mean Dice coefficient of 0.88 with a standard deviation of 0.1 across five-fold cross-validation on the segmentation task.
• The pre-trained model outperformed the randomly initialized U-net, which achieved a mean Dice coefficient of 0.84 with a standard deviation of 0.2.
• With only 10% of the training data, the pre-trained model achieved a Dice coefficient significantly higher than the baseline (p < 3.97×10⁻¹²).
• The performance gap between the two models was most pronounced at low data availability, demonstrating greater robustness with fewer segmentation labels.
• The paired t-test showed statistically significant differences (p < 0.05) at all tested data proportions, confirming the advantage of pre-training.
• The localization network achieved a mean squared error of 132.16 on the validation set and 206.15 on the test set.