Exploring Structure-Wise Uncertainty for 3D Medical Image Segmentation

TL;DR

This paper introduces a framework for measuring structure-wise uncertainty in 3D medical image segmentation, enhancing out-of-distribution detection and segmentation accuracy by focusing on local structures like tumors.

Contribution

The paper proposes a novel method to quantify structure-wise uncertainty, addressing a gap in existing voxel-wise and image-wise uncertainty measures for medical segmentation.

Findings

Structure-wise uncertainty improves segmentation accuracy.

The framework enhances out-of-distribution detection.

Validated on multiple tumor segmentation datasets.

Abstract

When applying a Deep Learning model to medical images, it is crucial to estimate the model uncertainty. Voxel-wise uncertainty is a useful visual marker for human experts and could be used to improve the model's voxel-wise output, such as segmentation. Moreover, uncertainty provides a solid foundation for out-of-distribution (OOD) detection, improving the model performance on the image-wise level. However, one of the frequent tasks in medical imaging is the segmentation of distinct, local structures such as tumors or lesions. Here, the structure-wise uncertainty allows more precise operations than image-wise and more semantic-aware than voxel-wise. The way to produce uncertainty for individual structures remains poorly explored. We propose a framework to measure the structure-wise uncertainty and evaluate the impact of OOD data on the model performance. Thus, we identify the best UE method to improve the segmentation quality. The proposed framework is tested on three datasets with the tumor segmentation task: LIDC-IDRI, LiTS, and a private one with multiple brain metastases cases.