MagicNet: Semi-Supervised Multi-Organ Segmentation via Magic-Cube Partition and Recovery

TL;DR

MagicNet introduces a novel semi-supervised multi-organ segmentation method using a partition-and-recovery strategy inspired by a magic cube, effectively leveraging anatomical priors to improve segmentation accuracy with limited labeled data.

Contribution

The paper proposes a new data augmentation approach based on partitioning CT volumes into N^3 cubes, guided by anatomical priors, to enhance semi-supervised multi-organ segmentation.

Findings

Outperforms state-of-the-art semi-supervised methods by +7% DSC on MACT dataset.

Effectively utilizes unlabeled data through the proposed augmentation strategy.

Improves segmentation of small organs with enhanced local attribute learning.

Abstract

We propose a novel teacher-student model for semi-supervised multi-organ segmentation. In teacher-student model, data augmentation is usually adopted on unlabeled data to regularize the consistent training between teacher and student. We start from a key perspective that fixed relative locations and variable sizes of different organs can provide distribution information where a multi-organ CT scan is drawn. Thus, we treat the prior anatomy as a strong tool to guide the data augmentation and reduce the mismatch between labeled and unlabeled images for semi-supervised learning. More specifically, we propose a data augmentation strategy based on partition-and-recovery N$^3$ cubes cross- and within- labeled and unlabeled images. Our strategy encourages unlabeled images to learn organ semantics in relative locations from the labeled images (cross-branch) and enhances the learning ability for small organs (within-branch). For within-branch, we further propose to refine the quality of pseudo labels by blending the learned representations from small cubes to incorporate local attributes. Our method is termed as MagicNet, since it treats the CT volume as a magic-cube and N$^3$-cube partition-and-recovery process matches with the rule of playing a magic-cube. Extensive experiments on two public CT multi-organ datasets demonstrate the effectiveness of MagicNet, and noticeably outperforms state-of-the-art semi-supervised medical image segmentation approaches, with +7% DSC improvement on MACT dataset with 10% labeled images. Code is available at https://github.com/DeepMed-Lab-ECNU/MagicNet.