DiffRect: Latent Diffusion Label Rectification for Semi-supervised Medical Image Segmentation

TL;DR

DiffRect introduces a novel semi-supervised medical image segmentation method that rectifies pseudo labels in latent space using diffusion models, significantly improving accuracy with limited labeled data.

Contribution

The paper proposes DiffRect, a latent diffusion-based framework with label calibration and rectification modules, enhancing pseudo label quality and segmentation performance in semi-supervised settings.

Findings

Achieves 82.40% Dice score on ACDC with 1% labeled data.

Outperforms previous methods by 4.60% in Dice score.

Rivals fully supervised performance on multiple datasets.

Abstract

Semi-supervised medical image segmentation aims to leverage limited annotated data and rich unlabeled data to perform accurate segmentation. However, existing semi-supervised methods are highly dependent on the quality of self-generated pseudo labels, which are prone to incorrect supervision and confirmation bias. Meanwhile, they are insufficient in capturing the label distributions in latent space and suffer from limited generalization to unlabeled data. To address these issues, we propose a Latent Diffusion Label Rectification Model (DiffRect) for semi-supervised medical image segmentation. DiffRect first utilizes a Label Context Calibration Module (LCC) to calibrate the biased relationship between classes by learning the category-wise correlation in pseudo labels, then apply Latent Feature Rectification Module (LFR) on the latent space to formulate and align the pseudo label distributions of different levels via latent diffusion. It utilizes a denoising network to learn the coarse to fine and fine to precise consecutive distribution transportations. We evaluate DiffRect on three public datasets: ACDC, MS-CMRSEG 2019, and Decathlon Prostate. Experimental results demonstrate the effectiveness of DiffRect, e.g. it achieves 82.40\% Dice score on ACDC with only 1\% labeled scan available, outperforms the previous state-of-the-art by 4.60\% in Dice, and even rivals fully supervised performance. Code is released at \url{https://github.com/CUHK-AIM-Group/DiffRect}.