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

Adway Kanhere, Pranav Kulkarni, Paul H. Yi, Vishwa S. Parekh; Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, 2024, pp. 2380-2387

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 this https://github.com/UM2ii/SegViz.