Table of Contents
Fetching ...

Uncertainty Driven Bottleneck Attention U-net for Organ at Risk Segmentation

Abdullah Nazib, Riad Hassan, Zahidul Islam, Clinton Fookes

TL;DR

The paper introduces an uncertainty-driven bottleneck attention (UDBA) mechanism within a dual-decoder U-net for organ-at-risk segmentation in CT images, leveraging decoder disagreement as a surrogate for uncertainty to refine bottleneck features. It also adds a CT intensity integrated regularization to better capture tissue contrast, including a matrix CTR variant to account for intra- and inter-class tissue contrast. Evaluations on SegThor and LCTSC datasets show consistent performance gains over strong baselines, with notable improvements in Dice and IoU across multiple organs. These contributions enhance segmentation accuracy in low-contrast CT scenarios, potentially improving precision in radiation treatment planning without requiring additional data acquisitions.

Abstract

Organ at risk (OAR) segmentation in computed tomography (CT) imagery is a difficult task for automated segmentation methods and can be crucial for downstream radiation treatment planning. U-net has become a de-facto standard for medical image segmentation and is frequently used as a common baseline in medical image segmentation tasks. In this paper, we propose a multiple decoder U-net architecture and use the segmentation disagreement between the decoders as attention to the bottleneck of the network for segmentation refinement. While feature correlation is considered as attention in most cases, in our case it is the uncertainty from the network used as attention. For accurate segmentation, we also proposed a CT intensity integrated regularization loss. Proposed regularisation helps model understand the intensity distribution of low contrast tissues. We tested our model on two publicly available OAR challenge datasets. We also conducted the ablation on each datasets with the proposed attention module and regularization loss. Experimental results demonstrate a clear accuracy improvement on both datasets.

Uncertainty Driven Bottleneck Attention U-net for Organ at Risk Segmentation

TL;DR

The paper introduces an uncertainty-driven bottleneck attention (UDBA) mechanism within a dual-decoder U-net for organ-at-risk segmentation in CT images, leveraging decoder disagreement as a surrogate for uncertainty to refine bottleneck features. It also adds a CT intensity integrated regularization to better capture tissue contrast, including a matrix CTR variant to account for intra- and inter-class tissue contrast. Evaluations on SegThor and LCTSC datasets show consistent performance gains over strong baselines, with notable improvements in Dice and IoU across multiple organs. These contributions enhance segmentation accuracy in low-contrast CT scenarios, potentially improving precision in radiation treatment planning without requiring additional data acquisitions.

Abstract

Organ at risk (OAR) segmentation in computed tomography (CT) imagery is a difficult task for automated segmentation methods and can be crucial for downstream radiation treatment planning. U-net has become a de-facto standard for medical image segmentation and is frequently used as a common baseline in medical image segmentation tasks. In this paper, we propose a multiple decoder U-net architecture and use the segmentation disagreement between the decoders as attention to the bottleneck of the network for segmentation refinement. While feature correlation is considered as attention in most cases, in our case it is the uncertainty from the network used as attention. For accurate segmentation, we also proposed a CT intensity integrated regularization loss. Proposed regularisation helps model understand the intensity distribution of low contrast tissues. We tested our model on two publicly available OAR challenge datasets. We also conducted the ablation on each datasets with the proposed attention module and regularization loss. Experimental results demonstrate a clear accuracy improvement on both datasets.
Paper Structure (13 sections, 7 equations, 2 figures, 4 tables)

This paper contains 13 sections, 7 equations, 2 figures, 4 tables.

Figures (2)

  • Figure 1: Overview of the proposed approach.A) The U-net with auxiliary decoder. B) The Uncertainty module, C) Attention module.
  • Figure 2: Sample segmentation images from different methods on SegThor (top two rows) and LCTSC dataset (bottom tow rows). The organs are Esophagus and Heart from both datasets. Red=Ground-truth, Green=Prediction.