Federated-Learning-Glaucoma

Federated Learning for Glaucoma Segmentation: Model Checkpoints

Overview

This repository contains trained model checkpoints from the research project: "A Federated Learning-based Optic Disc and Cup Segmentation Model for Glaucoma Monitoring in Color Fundus Photographs"

Key Information

• Task: Automated optic disc and cup segmentation for glaucoma assessment
• Architecture: Mask2Former with Swin Transformer backbone
• Pre-training: ADE20K semantic segmentation dataset
• Training Data: 5,550 color fundus photographs from 9 datasets across 7 countries
• Approach: Privacy-preserving federated learning with site-specific fine-tuning

Clinical Context

Glaucoma is a leading cause of irreversible blindness worldwide, affecting 3.54% of the population aged 40-80 and projected to impact 111.8 million people by 2040. A key indicator of glaucoma severity is the vertical cup-to-disc ratio (CDR), with ratios ≥0.6 suggestive of glaucoma.

This work addresses the need for accurate automated segmentation while preserving patient data privacy across multiple clinical sites, enabling HIPAA/GDPR-compliant multi-institutional collaboration.

Models Included

This repository contains 22 trained models organized into four categories:

Baseline Models

• Central Model (1 model): Trained on pooled multi-site data, representing upper bound performance
• Local Models (9 models): Site-specific models trained on individual datasets, representing lower bound performance

Federated Learning Models

• Pipeline 1 (1 model): Global Validation
• Pipeline 2 (1 model): Weighted Global Validation
• Pipeline 3 (1 model): Onsite Validation
• Pipeline 4 (9 models): Fine-Tuned Onsite Validation

Usage

Download Specific Model

from huggingface_hub import hf_hub_download

Download central model

model_path = hf_hub_download( repo_id="sud11111/Federated-Learning-Glaucoma", filename="models/central/best_model.pt" )

Download fine-tuned model for specific dataset

model_path = hf_hub_download( repo_id="sud11111/Federated-Learning-Glaucoma", filename="models/pipeline4/chaksu/best_model.pt" )

Download All Models

from huggingface_hub import snapshot_download

Download entire models directory

local_dir = snapshot_download( repo_id="sud11111/Federated-Learning-Glaucoma", allow_patterns="models/**" ) print(f"Models downloaded to: {local_dir}")

Load and Perform Inference

import torch from transformers import Mask2FormerForUniversalSegmentation, Mask2FormerImageProcessor from PIL import Image

Load preprocessor

processor = Mask2FormerImageProcessor.from_pretrained( "facebook/mask2former-swin-base-ade-semantic" )

Load model architecture

model = Mask2FormerForUniversalSegmentation.from_pretrained( "facebook/mask2former-swin-base-ade-semantic", num_labels=4 # background, unlabeled, optic disc, optic cup )

Load trained weights

model.load_state_dict(torch.load(model_path)) model.eval()

Perform inference on fundus image

image = Image.open("fundus_image.jpg") inputs = processor(images=image, return_tensors="pt")

with torch.no_grad(): outputs = model(**inputs)

Post-process segmentation

predicted_segmentation = processor.post_process_semantic_segmentation( outputs, target_sizes=[image.size[::-1]] )[0]

Datasets

Training was performed across 9 public datasets spanning 7 countries, comprising a total of 5,550 color fundus photographs from at least 917 patients:

Data Split: Each dataset was divided into training (80%), validation (10%), and testing (10%) subsets. For datasets with multiple expert annotations, the STAPLE (Simultaneous Truth and Performance Level Estimation) method was used to generate consensus segmentation labels.

Model Architecture

• Base Model: Mask2Former
• Backbone: Swin Transformer (Swin-Base)
• Pre-training: ADE20K semantic segmentation dataset
• Input Resolution: 512×512 pixels
• Output Classes: 4 (background, unlabeled, optic disc, optic cup)
• Optimizer: AdamW (learning rate: 2×10⁻⁵)
• Loss Function: Multi-class cross-entropy
• Early Stopping: Patience of 7 epochs/rounds

Training Configuration

Common Hyperparameters

• Batch size: 8
• Learning rate: 2×10⁻⁵
• Optimizer: AdamW with weight decay
• Maximum epochs: 100 (with early stopping)
• Early stopping patience: 7 epochs/rounds
• Input size: 512×512 pixels (normalized)

---