Glaucoma Segmentation

Master's thesis: neural networks that outline the optic disc and cup in eye images.

Dual-decoder U-Net: one encoder feeds two decoder branches that output the optic disc and optic cup masks

About the project

Glaucoma is diagnosed partly from the ratio between the optic cup and the optic disc. For my master's thesis I designed, implemented and evaluated two network architectures for finding them: a cascade, where a second network looks for the cup inside the disc found by the first, and a dual-decoder network with one shared encoder.

The models were trained on the ORIGA dataset and tested on DRISHTI-GS. The best one reached a Dice score of 96.8 % for the disc and 89.7 % for the cup.

Highlights

Cascade and dual-decoder architectures in PyTorch
Region-of-interest detection and training in polar coordinates
Explainability with Grad-CAM

Under the hood

Exploratory data analysis showed class imbalance and very different image quality, which shaped the preprocessing
Before settling on CenterNet I tried simpler ways of finding the optic disc: thresholding, an intensity-weighted centroid, the brightest spot and template matching
A CenterNet model first finds the region of interest around the optic disc, and the segmentation runs on that crop in polar coordinates
Both architectures were built on three U-Net variants (Residual Attention U-Net++, Refined U-Net 3+ with CBAM and Swin-Unet) and trained on polar as well as Cartesian images
The loss combines Dice and cross-entropy; in the dual-decoder network the cup counts five times more than the disc
Predictions are cleaned up afterwards: holes are filled, only the largest region is kept and an ellipse is fitted to each mask
The vertical cup-to-disc ratio from the predictions was compared with values measured by experts, and a threshold of 0.63 taken from the ROC curve sorts eyes into glaucoma and healthy
Probability maps, activation maps and Grad-CAM show what the networks look at
Command-line scripts train and evaluate the models and segment new images with any of the architectures

When: 2023 – 2024
Type: University
Supervision: Supervised by doc. RNDr. Silvester Czanner, PhD.
Versions: Python 3.10
PyTorch 2.0
Stack: Python
PyTorch
OpenCV
NumPy
pandas
scikit-learn
Albumentations
Matplotlib
Jupyter
Weights & Biases

Gallery

Processing steps from a raw eye image to the detected optic disc — Preprocessing: black edges are removed, the image is padded and resized, and CenterNet finds the optic disc to crop the region of interest.

Four eye images with the expert outline, the model prediction and both contours overlaid — Results of the dual-decoder model: input, ground truth, prediction, and the true (blue) and predicted (black) contours of the optic disc and cup.

Diagram of two U-Net models in a row: the first outputs the disc mask, which crops the input for the second, which outputs the cup mask — Cascade architecture: a first U-Net segments the optic disc, its mask crops the image, and a second U-Net finds the cup inside it. It reached a Dice score of 96.8 % for the disc and 87.1 % for the cup.

Diagram of a network with one shared encoder and two decoders, one for the disc and one for the cup — Dual-decoder architecture: one shared encoder, a decoder for the optic disc and one for the cup. It reached a Dice score of 96.8 % for the disc and 89.7 % for the cup.

Line charts of loss and Dice score over the training epochs — Training history of the dual-decoder model: loss, Dice score and other metrics on the training and validation sets.

Charts comparing the predicted and the expert cup-to-disc ratio — Vertical cup-to-disc ratio, the main sign of glaucoma, computed from the predictions and compared with values measured by experts.