PaperlayerView exact revision on arXiv

Source-linked AI summary

SAUNet: Shape Attentive U-Net for Interpretable Medical Image Segmentation

Jesse Sun, Fatemeh Darbehani, Mark Zaidi, Bo Wang

arXiv:2001.07645v3eess.IVcs.CV

TL;DR

Ventricular segmentation is time-consuming and variable, while conventional CNNs raise concerns about transparency and robustness. SAUNet combines parallel shape and texture streams with attention-based saliency, and reports state-of-the-art results on SUN09 and AC17.

  • Problem

    Ventricular segmentation remains challenging, and conventional CNNs are described as lacking robustness and interpretability.

    Source: §1 · pp. 1–2

  • Method

    SAUNet pairs a gated shape stream with a texture stream and uses decoder attention to produce saliency maps at multiple resolutions.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    SAUNet yields state-of-the-art results on SUN09 and AC17, with learned attention maps analyzed for interpretability.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors present SAUNet as a step toward clinically adoptable deep-learning methods and further work on interpretability.

    Source: §7 · p. 8

    The active-contour loss discussed as prior work models contour length and area rather than object shape, and the authors observed slow learning and poor convergence.

    Source: §2.4 · p. 4

Abstract

from arXiv · show

Medical image segmentation is a difficult but important task for many clinical operations such as cardiac bi-ventricular volume estimation. More recently, there has been a shift to utilizing deep learning and fully convolutional neural networks (CNNs) to perform image segmentation that has yielded state-of-the-art results in many public benchmark datasets. Despite the progress of deep learning in medical image segmentation, standard CNNs are still not fully adopted in clinical settings as they lack robustness and interpretability. Shapes are generally more meaningful features than solely textures of images, which are features regular CNNs learn, causing a lack of robustness. Likewise, previous works surrounding model interpretability have been focused on post hoc gradient-based saliency methods. However, gradient-based saliency methods typically require additional computations post hoc and have been shown to be unreliable for interpretability. Thus, we present a new architecture called Shape Attentive U-Net (SAUNet) which focuses on model interpretability and robustness. The proposed architecture attempts to address these limitations by the use of a secondary shape stream that captures rich shape-dependent information in parallel with the regular texture stream. Furthermore, we suggest multi-resolution saliency maps can be learned using our dual-attention decoder module which allows for multi-level interpretability and mitigates the need for additional computations post hoc. Our method also achieves state-of-the-art results on the two large public cardiac MRI image segmentation datasets of SUN09 and AC17.

1. Introduction

The introduction frames ventricular segmentation as clinically important but time-consuming and variable, while conventional CNNs raise concerns about interpretability and robustness. SAUNet combines shape learning and attention with texture features, and reports state-of-the-art results on SUN09 and AC17.

1. Introduction

  • Ventricular volume measurements contribute to cardiovascular disease diagnosis and prognosis, including conditions associated with abnormal cardiac output.The introduction gives examples including atherosclerosis, heart failure, and cardiomyopathy.Source: §1 · p. 1
  • Ventricular endocardium segmentation relies heavily on manual annotations that are time-consuming, ambiguously scoped, and subject to intra- and inter-observer variability.The challenge concerns left and right ventricular endocardium at end-systolic and diastolic timepoints.Source: §1 · p. 1
  • CNNs can be difficult to interpret and are strongly influenced by dense pixel values, which the authors describe as less robust features than object shapes.The authors argue that limited transparency and robustness hinder translation into clinical settings.Source: §1 · p. 2
  • SAUNet processes shape features in a secondary stream and uses decoder spatial and channel-wise attention to produce interpretability at each U-Net resolution.The shape stream's output is a shape attention map.Source: §1 · p. 2
  • The authors position shape learning as a built-in model feature rather than a new loss function, noting that earlier medical-segmentation approaches used new loss functions for shape information.They present the approach as addressing transparency and robustness concerns in clinical translation.Source: §1 · p. 2
  • SAUNet achieves state-of-the-art results on SUN09 and AC17, with attention maps supporting interpretation at multiple resolutions.The authors analyze spatial and shape attention maps to interpret learned features.Source: §1 · p. 2

2. Related Works

This section reviews U-Net and DenseNet architectures, attention mechanisms, and active-contour losses as related approaches to medical-image segmentation and shape learning.

2.1. U-Net

  • U-Net uses a contracting path for feature information and a symmetric expansive path for localization, with skip connections passing high-resolution information between similar-resolution stages.Decoder blocks concatenate aligned encoder features and upsample the resulting feature maps.Source: §2.1 · pp. 2–3

2.2. DenseNet

  • DenseNet dense blocks concatenate each layer's output with feature maps entering later layers, promoting feature reuse; transition blocks reduce feature depth and resolution.The review notes DenseNet-169 had better performance than ResNet-50 with approximately 50% fewer parameters.Source: §2.2 · p. 3

2.3. Attention in CNNs

  • Post hoc interpretability commonly uses gradients after training, whereas trainable attention learns to reweight feature importance; squeeze-and-excitation modules redistribute channel importance but lack an obvious interpretability method.The reviewed attention estimator produces spatial maps that can be visualized and used as masks for weakly supervised segmentation.Source: §2.3 · p. 3
  • The proposed attention approach combines channel-wise squeeze-and-excitation modules with Jetley et al.'s spatial attention [20], and the authors report strong results with interpretability.Related work notes that post hoc methods generally do not expose decisions at different forward-propagation stages.Source: §2.3 · p. 3

2.4. Active Contour Loss for Learning Shapes

  • Active contour loss incorporates area and size information, but considers contour length and inside/outside area rather than object shape; the authors observed slow learning and poor convergence.Chen et al. combined the loss with a U-Net with dense blocks and reported that it outperformed other mainstream losses, such as cross entropy.Source: §2.4 · pp. 3–4

3. Method

SAUNet pairs a U-Net-like texture stream with a gated shape stream that processes shape and boundary information, using a dual-attention decoder.

3. Method

  • SAUNet combines a U-Net texture stream with a gated shape stream to process shape and boundary information alongside dense pixel features.The texture encoder uses DenseNet-121 blocks, and the decoder uses the proposed dual-attention block.Source: §3 · p. 4

3.1. Gated Shape Stream

The gated shape stream refines boundary and shape information using texture-encoder features, then combines its shape features with texture features in the final texture-stream layers.

3.1. Gated Shape Stream

  • The gated shape stream refines relevant boundary and shape information using features processed by the texture-stream encoder.Gated convolution fuses shape-stream and texture information.Source: §3.1 · p. 4

3.1.1 Gated Convolutional Layer

  • The gated convolution computes a boundary attention map from shape- and texture-stream features, then uses it to gate shape features without pooling in the shape stream.Texture features are bilinearly interpolated when needed to match shape-feature dimensions; each residual block is a shape-stream layer.Source: §3.1.1 · p. 4

3.1.2 Output of Gated Shape Stream

  • The shape stream predicts class boundaries under deep supervision and concatenates its predicted shape maps with Canny edges before joining texture features.Its boundary loss sends gradients back to the texture-stream encoders.Source: §3.1.2 · p. 4

3.2. Dual Attention Decoder Block

The decoder combines skip-connection features with lower-resolution decoder features, then applies spatial and channel-wise attention paths to produce fused attention features.

3.2. Dual Attention Decoder Block

  • The dual-attention decoder adds spatial attention for interpretability and channel-wise attention for improved performance after the normalized convolution of concatenated decoder features.The decoder input combines encoder skip features with lower-resolution decoder features.Source: §3.2 · pp. 4–5

3.2.1 Spatial Attention Path

  • The spatial path uses two 1x1 convolutions and a sigmoid to produce a one-channel spatial attention map, then repeats it across channels for element-wise fusion.The first convolution reduces channels to C/2 and the second to one.Source: §3.2.1 · p. 5

3.2.2 Channel Attention Path

  • The channel-wise path uses squeeze-and-excitation to assign each skip-connection channel a coefficient in [0, 1] and scale that channel.Source: §3.2.2 · p. 5

3.2.3 Channel and Spatial Attention

  • The block fuses spatial and channel attention by multiplying the channel features by spatial attention plus one, allowing attention to amplify rather than zero features.N denotes the Hadamard product.Source: §3.2.3 · p. 5

3.3. Dual-Task Loss Function for Learning Shapes

The objective combines segmentation losses with a boundary-prediction loss, balancing pixelwise classification, region overlap, and shape learning.

3.3. Dual-Task Loss Function for Learning Shapes

  • The proposed objective combines segmentation loss with boundary loss from the shape stream to promote precise segmentation and shape learning.Source: §3.3 · p. 5, §3.3.3 · p. 6

3.3.1 Cross Entropy Loss

  • Cross-entropy loss averages pixelwise prediction error across the image domain and classes.The passage defines ground-truth labels as one-hot values and predictions as per-pixel probabilities.Source: §3.3.1 · p. 5

3.3.2 Dice Loss

  • Dice loss derives from set overlap, with the Dice coefficient maximized at 1 for identical sets and minimized at 0 for disjoint sets.The loss is defined as an inverted, differentiable coefficient.Source: §3.3.2 · p. 6

3.3.3 Dual Task Loss

  • The total loss weights cross-entropy, Dice, and edge losses; equal weights worked well, and using either segmentation loss alone did not make the model learn well and quickly.Cross-entropy targets pixel values, while Dice loss accounts for overlap.Source: §3.3.3 · p. 6

4. Experiments

Experiments evaluated SAUNet on SUN09 left-ventricle segmentation and AC17 bi-ventricular segmentation, using dataset-specific preprocessing and training protocols.

4. Experiments

  • Experiments used SUN09 and AC17, z-score normalization, and runtime rotations, flips, elastic deformations, and gamma shifts.Training ran on one NVIDIA Quadro RTX 5000 GPU with 16 GB of memory; code was publicly available on GitHub.Source: §4 · p. 6

4.1. SUN09 Left Ventricle Segmentation Dataset

  • SUN09 used separate endocardium and epicardium datasets of 260 and 135 2D slices, respectively, center-cropped to 128×128 pixels, with two models trained.Each model used RAdam and was trained for 120 epochs with batch size 4.Source: §4.1 · p. 6

4.2. AC17 Bi-Ventricular Segmentation Dataset

  • AC17 comprised 200 MRI volumes from 100 patients, resampled in-plane to 1.25 mm, center-cropped or padded to 256×256, and split 80:20 for training and validation.The model was trained for 180 epochs; results were submitted on the test set at the convergence epoch.Source: §4.2 · p. 6

5. Experimental Results

The experiments report test-set Dice results on SUN09 and AC17, ablate the shape stream on AC17, and assess cross-dataset robustness using SUN09-trained models on AC17.

5.1. SUN09

  • SUN09 test-set Dice scores are reported alongside the five prior works with the highest reported accuracies, without hyperparameter tuning for SAUNet.The evaluation uses Dice, consistent with other benchmarks and works.Source: §5 · p. 6, Table 1 · p. 7, §5.1 · p. 7

5.2. AC17 Segmentation

  • SAUNet yielded state-of-the-art Dice results among end-to-end models for all AC17 classes, with notably stronger right-ventricle performance than previous works.The authors suggest the shape stream’s learning of the right ventricle’s irregular shape may help explain that result.Source: Table 2 · p. 7, §5.2 · p. 7

5.3. Ablation Studies

  • On the AC17 validation split, adding the shape stream substantially increased right-ventricle Dice; test examples also visibly favored the model with the stream.The ablation compared models with and without the shape stream, trained with RAdam at learning rate 5E-4 and evaluated at convergence.Source: Table 3 · p. 7, Figure 2 · p. 7, §5.3 · p. 7

5.4. Robustness of Shape Stream

  • For SUN09-trained models tested on AC17, the shape-stream model’s mIoU percentage drop was 1.76% smaller than the model without it.Only the left ventricle was evaluated because SUN09 provides labels for that class alone.Source: Table 4 · p. 7, §5.4 · p. 7

6. Interpretability

SAUNet exposes shape and spatial attention maps for interpretation, with decoder-resolution maps highlighting different regions and a faster alternative to post hoc saliency generation.

Interpretability

  • SAUNet exposes spatial attention maps for decoder blocks and shape maps that indicate whether learned class shapes match the structures of interest.The figure includes spatial attention from the last two dual-attention decoder blocks, intermediate and final shape attention, and SmoothGrad maps.Source: §6 · p. 8
  • Higher-resolution decoder block 2 focuses strongly on the right ventricle, while lower-resolution block 3 distributes activation more broadly across image regions.The authors associate broader lower-resolution attention with processing more global information; shape-stream maps narrow onto the structures of interest.Source: §6 · p. 8
  • For all 384 validation images, the proposed method generated saliency maps in 20 seconds, versus 24 minutes with SmoothGrad; it also provides maps at different levels and resolutions.SmoothGrad requires 25–50 forward-backward passes per image; the proposed method generates all multi-resolution maps on the initial forward pass.Source: §6 · p. 8

7. Conclusion

The conclusion presents SAUNet as an inherently interpretable segmentation model that learns robust shape features and reports state-of-the-art results on SUN09 and AC17.

7. Conclusion

  • SAUNet is presented as an inherently interpretable medical image segmentation model.The conclusion names the architecture Shape Attentive U-Net.Source: §7 · p. 8
  • The model learns robust object-shape features through its gated shape stream and uses attention-based built-in saliency maps for interpretability.Source: §7 · p. 8
  • The authors report state-of-the-art results on the public bi-ventricular segmentation datasets SUN09 and AC17.They express hope that the work will support clinical adoptability and inspire further interpretability research.Source: §7 · p. 8

Related papers

Loading 2001.07645v3…