Towards to Robust and Generalized Medical Image Segmentation Framework

TL;DR

This paper introduces a two-stage transfer learning framework that enhances the robustness and generalization of medical image segmentation models, especially with limited training data, by leveraging domain-specific knowledge and auxiliary reconstruction.

Contribution

A novel two-stage transfer learning framework utilizing unsupervised pretraining and auxiliary reconstruction for improved medical image segmentation.

Findings

Demonstrates superior robustness to corruption

Achieves high generalization on unseen datasets

Effective with limited training data

Abstract

Deep learning-based computer-aided diagnosis is gradually deployed to review and analyze medical images. However, this paradigm is restricted in real-world clinical applications due to the poor robustness and generalization. The issue is more sinister with a lack of training data. In this paper, we address the challenge from the transfer learning point of view. Different from the common setting that transferring knowledge from the natural image domain to the medical image domain, we find the knowledge from the same domain further boosts the model robustness and generalization. Therefore, we propose a novel two-stage framework for robust generalized medical image segmentation. Firstly, an unsupervised tile-wise autoencoder pretraining architecture is proposed to learn local and global knowledge. Secondly, the downstream segmentation model coupled with an auxiliary reconstruction network is designed. The reconstruction branch encourages the model to capture more general semantic features. Experiments of lung segmentation on multi chest X-ray datasets are conducted. Comprehensive results demonstrate the superior robustness of the proposed framework to corruption and high generalization performance on unseen datasets, especially under the scenario of the limited training data.