Improved Vessel Segmentation with Symmetric Rotation-Equivariant U-Net

2025-01-24

AI Summary

Retinal vessel segmentation suffers from poor directional robustness and inconsistent predictions due to the lack of rotational/reflection equivariance in conventional CNNs. To address this, we propose the first integration of a lightweight Symmetric Rotational Equivariant Convolution (SRE-Conv) into the U-Net backbone—achieving strict rotational and reflectional equivariance while drastically reducing computational overhead. Unlike standard group-equivariant CNNs, our approach avoids excessive parameter counts and memory consumption. Evaluated on DRIVE and STARE benchmarks, it outperforms both vanilla U-Net and state-of-the-art equivariant models: parameter count reduced by over 40%, GPU memory footprint during inference lowered, and segmentation performance variance under arbitrary rotations and horizontal/vertical flips significantly diminished. These results demonstrate superior directional robustness and practical clinical applicability.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Non-convex Optimization

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Automated segmentation plays a pivotal role in medical image analysis and computer-assisted interventions. Despite the promising performance of existing methods based on convolutional neural networks (CNNs), they neglect useful equivariant properties for images, such as rotational and reflection equivariance. This limitation can decrease performance and lead to inconsistent predictions, especially in applications like vessel segmentation where explicit orientation is absent. While existing equivariant learning approaches attempt to mitigate these issues, they substantially increase learning cost, model size, or both. To overcome these challenges, we propose a novel application of an efficient symmetric rotation-equivariant (SRE) convolutional (SRE-Conv) kernel implementation to the U-Net architecture, to learn rotation and reflection-equivariant features, while also reducing the model size dramatically. We validate the effectiveness of our method through improved segmentation performance on retina vessel fundus imaging. Our proposed SRE U-Net not only significantly surpasses standard U-Net in handling rotated images, but also outperforms existing equivariant learning methods and does so with a reduced number of trainable parameters and smaller memory cost. The code is available at https://github.com/OnofreyLab/sre_conv_segm_isbi2025.

Problem

• Rotation and Reflection Invariance
• Vascular Segmentation
• Convolutional Neural Networks (CNN)

Innovation

• Symmetric Rotation Invariant Convolution
• U-Net Improvement
• Retinal Vessel Segmentation