Source-linked AI summary
OptiModNet: A UNet-Transformer Hybrid with Grouped-Query and Channel Attention for Optic Disc and Cup Segmentation
Soumili Ghosh, Debapriya Roy, Aryan Das, Bikash Santra
TL;DR
Reliable, efficient optic disc and cup segmentation remains challenging because glaucoma screening depends on accurate boundaries while existing models trade local detail, global context, and computational cost. OptiModNet combines a lightweight UNet-Transformer design with grouped-query and channel attention plus Aggregated Pyramid Loss, achieving state-of-the-art results on REFUGE2 and ORIGA with substantially lower complexity.
Problem
Accurate optic disc and cup segmentation is important for glaucoma screening, but image variability and boundary ambiguity make consistent, efficient segmentation challenging.
Method
OptiModNet extends UNetR with grouped-query attention in the encoder, channel attention in the decoder, and Aggregated Pyramid Loss for multi-stage supervision.
Results
99.45% disc Dice and 93.23% cup Dice on REFUGE2 surpassed existing approaches by over 2.5%, while ORIGA achieved the highest reported accuracy with nearly 96% fewer parameters than the most complex model.
Takeaways & Limitations
OptiModNet provides accurate optic disc and cup segmentation across REFUGE2 and ORIGA while substantially reducing model complexity.
Abstract
from arXiv · showhide
Precise segmentation of the optic disc and cup is critical for the early detection and diagnosis of glaucoma. However, achieving consistently high performance across datasets while maintaining low computational requirements remains a significant challenge. In glaucoma detection, low-computation methods are crucial for enabling rapid, large-scale screening and facilitating deployment in resource-limited clinical environments. While deep learning models such as UNets, Vision Transformers (ViTs), and Diffusion models have demonstrated strong segmentation performance but these methods often come with substantial computational overhead. UNets are efficient at capturing local features but are limited in modeling global contextual information. Conversely, ViTs excel at long-range dependency modeling but are computationally intensive. Hybrid architectures, such as UNetR, which combine transformer-based encoders with UNet-style decoders, have shown improved performance but while incurring additional complexity. Considering these, in this work, we propose OptiModNet, a light weight novel hybrid architecture tailored for optic disc and cup segmentation. The model integrates diverse attention mechanisms at multiple stages of the network to enhance both local and global feature representation. We include an Aggregated Pyramid Loss that supervises predictions at multiple decoder depths, to promote better gradient flow and structural consistency. We evaluate OptiModNet on the REFUGE2 dataset for both optic disc and cup segmentation tasks. Our method achieves state-of-the-art performance, exceeding existing approaches by over 2.5\%, while maintaining high efficiency with only 3.73 GFLOPs and 1.93M parameters. The code is available at https://github.com/SG1947/OptiModNet.
1 Introduction
Optic disc and cup segmentation supports glaucoma screening but remains difficult because of anatomical and imaging challenges. OptiModNet addresses these limitations with a lightweight hybrid architecture combining grouped-query attention, channel attention, and multi-stage supervision.
- Motivation: OD/OC segmentation enables cup-to-disc ratio computation, large-scale screening, and reduced inter-observer variability, but low contrast, peripapillary atrophy, and vessel occlusion make it challenging.These factors affect optic cup and disc boundary delineation in retinal fundus images.
- Motivation: UNet models capture local features but struggle with long-range dependencies, whereas ViTs model global relationships at higher computational cost and require large annotated datasets.The paper motivates a hybrid design to combine local detail extraction with global anatomical context.
- Proposed method: OptiModNet adapts a UNetR-based architecture to 2D fundus images, using grouped-query attention in the encoder and channel attention in the decoder.The design targets efficient global-context modeling and fine boundary-detail extraction.
- Proposed method: Aggregated Pyramid Loss supervises multiple decoder stages to improve gradient flow and fine-detail segmentation.Multi-stage supervision is intended to promote structural consistency across decoder depths.
- Contributions: The model combines UNet, Transformer, and attention mechanisms to improve OD/OC segmentation accuracy with fewer parameters.The stated contribution emphasizes complementary local, global, and channel-focused representations.
2 Literature Survey
The literature spans encoder–decoder CNNs, Vision Transformers, diffusion models, and hybrid architectures for medical image segmentation. In optic disc and cup segmentation, prior work emphasizes accurate boundary delineation for clinically important glaucoma assessment.
- Encoder–Decoder Architectures: UNet variants improve medical image segmentation through residual feature reuse, attention-based localization, automated configuration, and boundary-aware modules.Examples include ResUNet, ARU-GD, nnUNet, and BAT.
- Vision Transformers: Vision Transformers capture global contextual relationships through self-attention, while Swin Transformer reduces computational constraints using hierarchical shifted-window features.These methods address information beyond CNNs’ primarily local receptive fields.
- Diffusion-Based Approaches: Diffusion-based models such as EnsemDiff, SegDiff, MedSegDiff, and MedSegDiff-V2 generate structured, sharper, and more consistent segmentation outputs but often increase computational cost.The passage also notes that combining Vision Transformers, UNet, and attention methods supports medical imaging.
- Optic Disc and Cup Segmentation: Optic disc and cup segmentation research prioritizes accurate boundary delineation because it is crucial for computing the CDR, an important biomarker for glaucoma diagnosis.Earlier studies adapted UNet for fundus images and used region-based convolutional networks for simultaneous optic disc and cup segmentation.
3 Method
OptiModNet combines a transformer encoder with a convolutional decoder, using Grouped Query Attention, channel attention, and multi-stage supervision to balance contextual modeling with efficient segmentation. Its decoder progressively reconstructs predictions through skip connections, feature recalibration, and convolutional refinement.
- OptiModNet replaces conventional multi-head self-attention with Grouped Query Attention, splitting queries into groups that share key-value pairs to reduce computational load while modeling global dependencies.The encoder stacks transformer blocks containing grouped query attention, MLP refinement, and residual connections.
- The encoder processes patch and positional embeddings through 12 transformer blocks to learn long-range interactions and contextual information.Learnable positional embeddings are added to projected image patches before transformer encoding.
- The decoder progressively upsamples features, integrates multi-resolution encoder skip connections, applies channel attention, and refines representations with convolution, batch normalization, and ReLU.Each upsampling stage doubles spatial resolution and combines the upsampled map with a corresponding encoder feature.
- Channel attention recalibrates the combined decoder features using global average pooling and learned channel weights, emphasizing informative channels while reducing less relevant ones.The recalibrated features are multiplied element-wise by reshaped attention weights before convolutional refinement.
- Aggregated Pyramid Loss supervises multiple output branches through a weighted sum, using λ = {0.2, 0.3, 0.5} for three intermediate outputs and λ = 1 for the final output.The weighting prioritizes the final prediction while retaining supervision at intermediate stages.
4 Experiments
Experiments show that OptiModNet achieves leading segmentation performance on REFUGE2 and ORIGA while using few parameters, with rapid convergence, strong boundary alignment, and meaningful contributions from its attention modules.
- Benchmark comparisons: 98.92% IoU and 99.31% Dice are reported for OptiModNet on ORIGA three-class segmentation, improving over Wang et al. by 1.32% and 0.52%.The three classes are optic cup, optic disc, and background.
- Benchmark comparisons: OptiModNet gains 7.34% and 6.53% IoU over U-Net and Attention U-Net, alongside Dice gains of 6.55% and 6.92%, respectively.The passage also states that the method compares favorably with Swin-UNet, TransUNet, and traditional convolutional baselines.
- Training behavior: Over 50 epochs on REFUGE2 optic-disc segmentation, training and validation Dice and IoU improve rapidly while loss decreases, indicating convergence.The curves show strong alignment with the ground truth during training and validation.
- Qualitative analysis: Predicted REFUGE2 optic-disc and cup masks generally align closely with ground-truth boundaries, producing smooth contours despite minor deviations from illumination variation or vessel occlusion.The qualitative results emphasize localization and shape preservation, with less distinct disc margins causing some boundary errors.
- Efficiency analysis: 1.9M parameters make OptiModNet the smallest compared model, representing 96% parameter efficiency relative to MedSegDiff v2.The comparison evaluates parameter counts and computational complexity in Gflops.
- Ablation study: The ablation study reports incremental score improvements as attention modules are added, while separately evaluating λ values for the proposed loss on REFUGE2 disc segmentation.A visual comparison with UNETR 2D also reports improved boundary adherence for OptiModNet.
5 Conclusion and Future Work
OptiModNet extends UNetR for optic disc and cup segmentation by integrating Grouped Query Attention and Aggregated Pyramid Loss. It achieves state-of-the-art REFUGE2 performance, with Dice scores of 99.45% for disc and 93.23% for cup, surpassing existing approaches by over 2.5%.
- Conclusion: OptiModNet extends the UNetR architecture for optic disc and cup segmentation.The framework integrates Grouped Query Attention and a newly introduced Aggregated Pyramid Loss function.
- Conclusion: 99.45% Dice for disc and 93.23% Dice for cup are achieved on REFUGE2.These results are reported as state-of-the-art performance on the REFUGE2 dataset.
- Conclusion: Over 2.5% improvement over existing approaches is reported on REFUGE2.The comparison is stated for the reported disc and cup segmentation performance.