# Double-mix pseudo-label framework: enhancing semi-supervised segmentation on category-imbalanced CT volumes

## TL;DR

This paper introduces a new method to improve CT image segmentation with limited labeled data by focusing on challenging categories.

## Contribution

The novel Double-Mix Pseudo-label Framework and confidence-difficulty weight allocation method improve semi-supervised segmentation.

## Key findings

- The method achieved a 5.1% improvement in Dice score for CHD dataset with 5% labeled data.
- It achieved a 7.0% improvement in Dice score for BTCV dataset with 40% labeled data.
- The approach enhances segmentation of difficult categories in CT volumes.

## Abstract

Deep-learning-based supervised CT segmentation relies on fully and densely labeled data, the labeling process of which is time-consuming. In this study, our proposed method aims to improve segmentation performance on CT volumes with limited annotated data by considering category-wise difficulties and distribution.

We propose a novel confidence-difficulty weight (CDifW) allocation method that considers confidence levels, balancing the training across different categories, influencing the loss function and volume-mixing process for pseudo-label generation. Additionally, we introduce a novel Double-Mix Pseudo-label Framework (DMPF), which strategically selects categories for image blending based on the distribution of voxel-counts per category and the weight of segmentation difficulty. DMPF is designed to enhance the segmentation performance of categories that are challenging to segment.

Our approach was tested on two commonly used datasets: a Congenital Heart Disease (CHD) dataset and a Beyond-the-Cranial-Vault (BTCV) Abdomen dataset. Compared to the SOTA methods, our approach achieved an improvement of 5.1% and 7.0% in Dice score for the segmentation of difficult-to-segment categories on 5% of the labeled data in CHD and 40% of the labeled data in BTCV, respectively.

Our method improves segmentation performance in difficult categories within CT volumes by category-wise weights and weight-based mixture augmentation. Our method was validated across multiple datasets and is significant for advancing semi-supervised segmentation tasks in health care. The code is available at https://github.com/MoriLabNU/Double-Mix.