MO-CTranS: A unified multi-organ segmentation model learning from multiple heterogeneously labelled datasets
Zhendi Gong, Susan Francis, Eleanor Cox, Stamatios N. Sotiropoulos, Dorothee P. Auer, Guoping Qiu, Andrew P. French, Xin Chen
TL;DR
This work tackles multi-organ segmentation when available data come from multiple small, heterogeneously labelled datasets. It introduces MO-CTranS, a unified CNN-Transformer framework with pyramid feature fusion and a task-specific token in the deepest Transformer level to differentiate label discrepancies across datasets. The approach outperforms baseline and state-of-the-art methods on abdominal MRI datasets across different views, while maintaining a lighter architecture and better resilience to label conflicts and data imbalance. The results demonstrate the practical impact of joint learning from diverse datasets for robust, cross-task organ segmentation, with potential for continual learning as more data become available.
Abstract
Multi-organ segmentation holds paramount significance in many clinical tasks. In practice, compared to large fully annotated datasets, multiple small datasets are often more accessible and organs are not labelled consistently. Normally, an individual model is trained for each of these datasets, which is not an effective way of using data for model learning. It remains challenging to train a single model that can robustly learn from several partially labelled datasets due to label conflict and data imbalance problems. We propose MO-CTranS: a single model that can overcome such problems. MO-CTranS contains a CNN-based encoder and a Transformer-based decoder, which are connected in a multi-resolution manner. Task-specific tokens are introduced in the decoder to help differentiate label discrepancies. Our method was evaluated and compared to several baseline models and state-of-the-art (SOTA) solutions on abdominal MRI datasets that were acquired in different views (i.e. axial and coronal) and annotated for different organs (i.e. liver, kidney, spleen). Our method achieved better performance (most were statistically significant) than the compared methods. Github link: https://github.com/naisops/MO-CTranS.