Retina Vessel Segmentation U-Net
This model is a U-Net architecture trained for binary segmentation of retinal vessels from fundus images. It was trained on a combined dataset (DRIVE, HRF, CHASE_DB1, and STARE).
Visual Results
Below is a comparison showing the original fundus image, the manual ground truth annotation, and the mask predicted by this U-Net model.
Model Architecture
The architectural foundation for the segmentation task is the U-Net model, implemented in PyTorch. This architecture is renowned for its efficacy in biomedical image segmentation. The model features a symmetric encoder-decoder structure with skip connections that bridge corresponding convolutional blocks.
Architectural Details
- Encoder (Contracting Path): The encoder consists of four blocks, each containing two 3×3 convolutional layers followed by a ReLU activation. After each block, a 2×2 max-pooling operation with a stride of 2 is applied for downsampling. The feature map depth progressively increases from an initial 128 channels to 256, 512, and finally 1024 channels.
- Bottleneck: At the lowest point of the U, a bottleneck block with two convolutional layers transforms the feature maps to a depth of 2048 channels.
- Decoder (Expansive Path): The decoder symmetrically mirrors the encoder, using 2×2 transposed convolutions to upsample the feature maps. The upsampled features are concatenated with the corresponding feature maps from the encoder via skip connections. This critical step allows the network to recover fine-grained spatial information lost during downsampling.
- Output Layer: The final layer consists of a 1×1 convolution that maps the 128-channel feature map to a single-channel output, followed by a Sigmoid activation function to produce a pixel-wise probability map between 0 and 1.
Training Protocol
The model was trained on a device equipped with a CUDA-enabled GPU. A custom PyTorch Dataset class was implemented to handle the loading and preprocessing of image-mask pairs.
- Configuration: Batch size of 4, trained for 250 epochs.
- Loss Function: A composite loss function, combining Binary Cross-Entropy (BCE) with Dice Loss in equal measure ($L=0.5 \times L_{BCE} + 0.5 \times L_{Dice}$). This hybrid approach leverages pixel-wise accuracy while maintaining robustness against class imbalance.
- Optimizer: Adam optimizer with a learning rate of $1 \times 10^{-4}$.
Performance Metrics (Test Set)
- Test Accuracy: 0.9665
- Test Precision: 0.8124
- Test Recall: 0.7539
- Test F1 Score: 0.7820
- AUC: 0.98
Intended Use
This model is intended for research purposes in medical image analysis, specifically for automated retinal vessel mapping.