Privacy-Preserving Collaboration for Multi-Organ Segmentation via Federated Learning from Sites with Partial Labels

Adway U. Kanhere, Pranav Kulkarni, Paul H. Yi, Vishwa S. Parekh

Abstract: Manual annotation of 3D medical images is expensive and time-consuming, resulting in datasets focused on segmenting individual organs. This leads to training several specialized models that limit clinical translational utility. To that end, we developed SegViz, a federated learning (FL) framework to aggregate knowledge from heterogeneous datasets with partial annotations into a single multi-organ segmentation model. SegViz uses collaborative 3D-U-Nets, with selective weight synchronization across distributed sites, to consolidate knowledge by averaging shared representation weights while isolating task-specific heads during synchronization. SegViz was compared to conventional FL using FedAvg, single-organ baseline models, and a single centralized model trained using data aggregated from all sites. Four partially annotated datasets were used in this study: Spleen MSD, Liver MSD, Pancreas MSD, and the Kidney Tumor Segmentation dataset. All approaches were evaluated using the independent BTCV dataset for segmentation of liver, spleen, pancreas, and kidneys using the dice similarity metric. Extensive experiments across the two-, three- and four-client FL setups with each client holding a dataset with single-organ annotations demonstrated the effectiveness of SegViz for collaborative multi-task segmentation from distributed sites with partial labels. All our implementations and code are available at https://github.com/UM2ii/SegViz.