AI Research Portal

This paper introduces AGENet, a novel framework for few-shot medical image segmentation that leverages edge-aware geodesic distance learning to incorporate spatial relationships and anatomical boundaries. The method aims to improve precise boundary delineation, which is a challenge for existing few-shot segmentation approaches, especially in medical imaging where data is scarce. AGENet achieves this by using computationally lightweight geometric modeling to guide prototype extraction, leading to improved performance across diverse medical imaging datasets.

The paper addresses several key challenges in few-shot medical image segmentation:

Suboptimal Boundary Delineation: Existing few-shot segmentation methods struggle to accurately capture the precise boundaries of anatomical structures, particularly when anatomically similar regions appear without sufficient spatial context.
Limited Spatial Modeling: Current prototype aggregation methods treat all spatial locations within anatomical structures uniformly, failing to account for the anisotropic nature of medical structures and the importance of boundary information.

An edge-aware geodesic distance learning module that incorporates anatomical boundaries through iterative Fast Marching refinement and adaptive parameter learning.
An adaptive prototype extraction framework that captures both global anatomical structure and local boundary details using spatially-weighted aggregation and adaptive grid sampling.
Comprehensive evaluation and analysis across nine diverse medical imaging datasets, demonstrating significant improvements over state-of-the-art methods.

AGENet is designed to enhance prototype-based few-shot segmentation by incorporating spatial anatomical information. The framework consists of two main components:

Edge-Aware Geodesic Distance Learning: This module transforms uniform spatial treatment into anatomically-informed distance fields. It begins with Euclidean Distance Transform (EDT) initialization, followed by iterative Fast Marching refinement to incorporate edge-aware constraints. Adaptive parameter learning is used to tailor the process to specific image characteristics, determining an optimal center-boundary balance. The speed function is formulated to slow propagation at detected boundaries, ensuring geodesic paths respect anatomical structures.
- EDT Initialization: G = EDT(M) initializes the distance field.
- Edge Detection: Sobel operators are used to detect anatomical boundaries: ||∇M||² = (Kx * M)² + (Ky * M)² + ε.
- Speed Function: Fi,j = 1 / (1 + β||∇G||²) controls geodesic path propagation, where β is a learnable parameter.
- Iterative Refinement: The distance field is refined using an update term U and iterative updates G(t) based on the Eikonal equation.
- Adaptive Parameter Learning: Geometric centrality (Cg) and boundary strength (bs) are used to compute an adaptive weighting parameter θ. A boundary sensitivity parameter β is computed based on speed function variance (sv). These parameters guide the weighting function w(G) = θ tanh(β · G) + (1 − θ) · (1 – tanh(β · G)).
Adaptive Grid Prototype Extraction: This component constructs spatially-informed prototypes that capture local and global anatomical patterns.
- Global Prototype Extraction: Spatially-weighted global prototypes (Pglobal) are computed using the adaptive weighting function.
- Adaptive Grid Generation and Prototype Integration: An adaptive prototype grid is created with concentrated sampling density near boundaries. Each grid cell prototype (Pgrid) aggregates local features within its neighborhood. The final representation (Punified) combines grid and global prototypes: Punified = {Pgrid} ∪ {Pglobal}.

The overall loss function combines segmentation loss (Lseg), edge-aware loss (Ledge), and alignment loss (Lalign): Ltotal = Lseg + Ledge + Lalign.

Architecture: The network uses a pre-trained ResNet-101 based DeepLab encoder. Input images are scaled to 3x256x256, resulting in feature maps of 256x32x32. The adaptive prototype grid operates at 8x8 resolution.
Key Equations:
- EDT Initialization: G = EDT(M)
- Edge Detection: ||∇M||² = (Kx * M)² + (Ky * M)² + ε
- Speed Function: Fi,j = 1 / (1 + β||∇G||²)
- Adaptive Parameter θ: θ = Cg / (Cg + bs + ε)
- Adaptive Parameter β: β = βbase + λ / (sv + ε)
- Weighting Function: w(G) = θ tanh(β · G) + (1 − θ) · (1 – tanh(β · G))
- Global Prototype: Pglobal = Σi,j Mi,j · W(Gi,j) · Fi,j / Σi,j Mi,j · W(Gi,j) + ε
- Grid Prototype: Pgrid = Σ(i,j)∈Ng,h Pi,j Wi,j · Mi,j · Fi,j / Σ(i,j)∈Ng,h Pi,j Wi,j · Mi,j + ε
- Segmentation Loss: Lseg = LCE + LDice
- Edge-Aware Loss: Ledge = (1 / HW) Σi,j BCE(Epred, Etarget)
- Alignment Loss: Lalign = (1 / NW) Σn=1 Σw=1 LCE(predsupport, labelpseudo)
- Total Loss: Ltotal = Lseg + Ledge + Lalign
Loss/Objective: The total loss comprises segmentation loss (Lseg), edge-aware loss (Ledge), and alignment loss (Lalign). Lseg combines cross-entropy and Dice loss. Ledge emphasizes boundary accuracy using Sobel edge detection. Lalign enforces bidirectional consistency between support and query representations.
Optimization: SGD optimizer with an initial learning rate of 0.005 and a multi-step scheduler reducing by 0.95.
Augmentation: Not stated in the paper.
Data Handling: Input images are scaled to 3x256x256.
Implementation Details: Implemented in PyTorch. Trained for 100,000 iterations.

Datasets: Experiments were conducted on nine datasets: ACDC, ABD-MRI, ABD-CT, PH2 (Atypical Nevus), PH2 (Melanomas), CVC-300, CHASE-DB1, Drishti-GS, and DCA1.
Metrics: DICE Similarity Coefficient (%) and 95th percentile Hausdorff Distance (HD95) in mm.
Baselines: PANet [37], SE-Net [25], SSL-ALPNet [19], Q-Net [28], CAT-Net [13], DSPNet [32], and RPT [39].
Numeric Results:
- Overall: Achieved 79.56% (1-shot) and 81.67% (5-shot) mean DICE scores.
- HD95: Achieved 11.16mm (1-shot) and 8.39mm (5-shot) Hausdorff Distance.
- Specific Datasets (5-shot): Achieved 97.39% DICE on CHASE-DB1 and 62.93% DICE on DCA1.
- Organ-Specific (5-shot): Achieved 91.52% DICE for right kidney in ABD-MRI and 82.63% DICE for left ventricle in ACDC.
Runtime/Training Cost: Trained on NVIDIA RTX 4090 GPU for approximately 4 hours per fold. Inference time is 11.17 ms.

No explicit paper-reported limitations were found.

Method evidence: "Our key insight is that medical structures follow predictable geometric patterns that can guide prototype extraction even with limited training data. Unlike methods relying on complex architectural components or heavy neural networks, our approach leverages computationally lightweight geometric modeling." (Abstract)
Result evidence: "Experimental results demonstrate substantial improvements over state-of-the-art methods, achieving 79.56% (1-shot) and 81.67% (5-shot) mean Dice scores with 11.16mm and 8.39mm Hausdorff Distance, while maintaining computational efficiency competitive with existing approaches." (Section 4.1)
Limitation evidence: Not stated in the paper.

AGENet: Adaptive Edge-aware Geodesic Distance Learning for Few-Shot Medical Image Segmentation

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures