Source-linked AI summary
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
Changlu Guo, Márton Szemenyei, Yugen Yi, Wenle Wang, Buer Chen, Changqi Fan
TL;DR
Retinal vessel segmentation is important for analyzing vascular changes associated with diseases such as diabetes and hypertension, but small datasets and difficult vessel-background distinctions challenge U-Net-based methods. SA-UNet combines a lightweight U-Net, structured dropout convolutional blocks, and spatial attention, achieving state-of-the-art performance on DRIVE and CHASE_DB1.
Problem
Retinal vessel segmentation is important for vascular analysis, while small training datasets and complex, low-contrast vessel structures make accurate segmentation difficult.
Method
SA-UNet replaces U-Net convolutional blocks with structured dropout blocks and adds spatial attention to refine important retinal vessel features.
Results
SA-UNet achieves state-of-the-art retinal vessel segmentation performance on both DRIVE and CHASE_DB1.
Takeaways & Limitations
The reported results support structured dropout and spatial attention as effective components for lightweight retinal vessel segmentation.
Abstract
from arXiv · showhide
The precise segmentation of retinal blood vessels is of great significance for early diagnosis of eye-related diseases such as diabetes and hypertension. In this work, we propose a lightweight network named Spatial Attention U-Net (SA-UNet) that does not require thousands of annotated training samples and can be utilized in a data augmentation manner to use the available annotated samples more efficiently. SA-UNet introduces a spatial attention module which infers the attention map along the spatial dimension, and multiplies the attention map by the input feature map for adaptive feature refinement. In addition, the proposed network employs structured dropout convolutional blocks instead of the original convolutional blocks of U-Net to prevent the network from overfitting. We evaluate SA-UNet based on two benchmark retinal datasets: the Vascular Extraction (DRIVE) dataset and the Child Heart and Health Study (CHASE_DB1) dataset. The results show that the proposed SA-UNet achieves state-of-the-art performance on both datasets.The implementation and the trained networks are available on Github1.
I. INTRODUCTION
Retinal vessel segmentation supports vascular analysis for diseases including diabetes and hypertension but remains difficult because vessels are small, connected, low-contrast, and affected by noise. SA-UNet addresses these challenges with spatial attention and structured dropout convolutional blocks in a lightweight U-Net.
- Retinal vessel segmentation provides morphological information about the vascular tree for quantitative fundus-image analysis.
- Small fragile vessels, connected structures, low vessel-background contrast, uneven lighting, and noise make segmentation challenging.
- U-Net variants perform well but can become more complex and less interpretable.
- SA-UNet introduces spatial attention into U-Net and uses a lightweight architecture to address retinal vessel segmentation challenges.
- Its structured dropout convolutional block combines DropBlock and batch normalization, while spatial attention enhances vascular features and suppresses unimportant features.
II. METHODOLOGY
SA-UNet uses a U-shaped encoder-decoder with structured dropout convolutional blocks and a spatial attention module between the encoder and decoder. Skip connections combine encoder features with decoder representations during upsampling.
- The network has a U-shaped encoder-decoder, with structured dropout convolutional blocks throughout the encoder and decoder.
- Each encoder step applies convolution, DropBlock, batch normalization, ReLU, and 2×2 max pooling, doubling feature channels during downsampling.
- Each decoder step uses 2×2 transposed convolution, halves feature channels, concatenates the corresponding encoder feature map, and applies a structured dropout block.
- The spatial attention module is placed between the encoder and decoder.
B. Structured Dropout Convolutional Block
The structured dropout convolutional block replaces U-Net’s original block to reduce overfitting, while batch normalization accelerates convergence. Spatial attention further refines features by modeling spatial relationships.
- Structured Dropout Convolutional Block: Original U-Net and lightweight U-Net training still show overfitting despite data augmentation, motivating DropBlock regularization.
- Structured Dropout Convolutional Block: DropBlock discards contiguous feature-map regions rather than independent random units.
- Structured Dropout Convolutional Block: The structured dropout block orders each convolutional layer with DropBlock, batch normalization, and ReLU activation.
- Structured Dropout Convolutional Block: The Backbone uses 18 convolutional layers instead of the original U-Net’s 23 and is reported to solve overfitting while accelerating convergence.
- Spatial Attention Module: Spatial attention pools features by channel-wise maximum and average operations, concatenates the descriptors, and generates a spatial attention map.
- Spatial Attention Module: The attention-map generation uses a convolution with a 7×7 kernel followed by a Sigmoid activation.
A. Datasets
SA-UNet is evaluated on the DRIVE and CHASE_DB1 retinal fundus datasets. Because the original image sizes are unsuitable for the network, images are padded for processing and cropped back during evaluation.
- The evaluation uses the public DRIVE and CHASE_DB1 retinal fundus image datasets.
- Images are zero-padded to fit the network and cropped to their initial size during evaluation.
- Four augmentation methods expand each dataset from 20 original training images to 256 images.
B. Evaluation Metrics
The study evaluates retinal vessel segmentation using confusion-matrix-derived metrics and AUC, with MCC emphasized for imbalanced binary classification.
- Sensitivity, specificity, F1, and accuracy are computed from pixel-level TP, FP, FN, and TN comparisons against ground truth.
- AUC measures segmentation performance, with AUC = 1 indicating perfect segmentation.
C. Implementation Details
SA-UNet is trained from scratch with augmented data using Adam and binary cross entropy, with dataset-specific batch sizes and DropBlock rates.
- SA-UNet is trained from scratch on augmented training sets using the Adam optimizer and binary cross entropy loss.
- The DRIVE configuration uses batch size 8 and DropBlock dropout rate 0.18, whereas CHASE DB1 uses batch size 4 and rate 0.13.
- Experiments use Keras with TensorFlow backend and an NVIDIA TITAN XP GPU with 12 Gigabyte memory.
IV. RESULTS
Ablation experiments on DRIVE and CHASE_DB1 examine spatial attention, structured dropout, and batch normalization, while parameter counts assess model size.
- Ablation Experiments: The ablations compare U-Net, U-Net + SA, SD-Unet, Backbone, and SA-UNet across DRIVE and CHASE_DB1.Backbone denotes SD-Unet + BN, while SA-UNet denotes Backbone + SA.
- Ablation Experiments: With only 98 parameters added, U-Net + SA performs better than U-Net, supporting the effectiveness of spatial attention.
- Ablation Experiments: Structured dropout improves Backbone over U-Net on DRIVE and CHASE_DB1 across ACC, AUC, F1, and MCC.Reported gains are 0.28% / 0.22%, 0.73% / 0.59%, 2.42% / 2.48%, and 2.48% / 2.64%, respectively.
- Ablation Experiments: Batch normalization improves Backbone over SD-Unet, although it slightly increases the parameter count.
- Ablation Experiments: SA-UNet achieves the best performance on most metrics while using fewer parameters than compared larger models.
- Ablation Experiments: Visual examples show SA-UNet retaining more accurate small and border-vessel structures than U-Net and SD-Unet.
B. Comparisons with state-of-the-art methods
Comparisons with existing retinal vessel segmentation methods show that SA-UNet performs best on both DRIVE and CHASE_DB1, with strong reported sensitivity, accuracy, and AUC.
- SA-UNet achieves the best performance on both DRIVE and CHASE_DB1 among the compared state-of-the-art methods.
- Sensitivity is 0.8212 / 0.8573, accuracy is 0.9698 / 0.9755, and AUC is 0.9864 / 0.9905.The passage reports these values for DRIVE / CHASE_DB1, respectively.
- Specificity is comparable with other methods, while SA-UNet uses much fewer parameters than AG-Net.
- On visual examples, SA-UNet has better segmentation at intersections of small blood vessels than AG-Net.
V. CONCLUSION
SA-UNet adapts U-Net for small retinal vessel datasets by combining structured dropout convolutional blocks with spatial attention. Evaluations on DRIVE and CHASE_DB1 support its effectiveness and state-of-the-art performance among compared retinal vessel segmentation methods.
- Method: SA-UNet replaces U-Net convolutional blocks with structured dropout blocks integrating DropBlock and batch normalization to address overfitting on small datasets.The approach uses data augmentation and a lightweight U-Net backbone because retinal datasets typically contain few annotated samples.
- Method: The network inserts a spatial attention module between the encoder and decoder to emphasize important features and suppress unnecessary ones.This targets low-contrast vessel regions, especially vessel edges and small vessels.
- Evaluation: SA-UNet was evaluated on the publicly available DRIVE and CHASE_DB1 retinal fundus image datasets.
- Results: The experiments demonstrate that structured dropout and spatial attention are effective, while comparisons with other state-of-the-art methods show state-of-the-art performance for the lightweight SA-UNet.