Exploring Structure-Wise Uncertainty for 3D Medical Image Segmentation

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.

Experimental Setup

1.1 Preprocessing

Here, we describe data preparation steps including datasets splits, normalization, and interpolation.

Mets data is randomly split into train (1140 images) and test (414 images) sets. We interpolate the images to have $1\,\text {mm} \times 1\,\text {mm} \times 1$ mm spacing.

LIDC data is randomly split into train (812 images) and test (204 images) sets. We clip image intensities between −1350 and 350 Hounsfield units (HU)—the standard lung window. We interpolate images to have 1 mm $\times$ 1 mm $\times$ 1.5 mm spacing.

LiTS is presented as two subsets, so we use the first as a test (28 images) and the second, excluding cases with empty tumor masks, as a train (90 images) set. The images are cropped to the provided liver masks. The intensities are clipped to the $[-150, 250]$ HU interval—the standard liver window. Finally, we interpolate images to have 0.77 mm $\times$ 0.77 mm $\times$ 1 mm spacing.

LiTS-mod is obtained by random changes of the reconstruction kernel to be extremely soft ($a=-0.7, b=0.5$) or sharp ($a=30, b=3$) using the implementation and notations of [26], and addition of “metal” artifacts (ball of radius 5 and 3000 HU) by substituting the parts of sinogram projection, as in [25].

Before passing through the network, we scale image intensities in [0, 1].

1.2 Training Setup

Although using cross-entropy loss has theoretical justifications of encouraging better calibrated predictions [16], models trained with this loss function fail in our segmentation task. For that reason we use Dice Loss [28] and its modifications in our experiments. Thus, uncertainty estimates might be shifted in such tasks, and experimental evaluation, as in our study, becomes even more relevant. All models are trained in a patch-based manner: patches are sampled randomly so that they contain structures. We use SGD optimizer with Nesterov momentum of 0.9 and $10^{-3}$ initial learning rate, which is decreased to $10^{-4}$ after $80\%$ of epochs. For LiTS and Mets segmentation the model is trained for 100 epochs (100 iterations per epoch, batch size 20), while for LIDC segmentation there are 30 epochs (1000 iterations per epoch, batch size 2).