tomalmog/oct-retinal-classifier

Retinal OCT Disease Classifier

A deep learning model for classifying retinal diseases from Optical Coherence Tomography (OCT) images. The model distinguishes between four conditions:

• CNV - Choroidal Neovascularization
• DME - Diabetic Macular Edema
• DRUSEN - Early AMD (Age-related Macular Degeneration)
• NORMAL - Healthy Retina

Model Description

This model uses EfficientNet-B3 as the backbone with a custom classification head, trained on the Kermany2018 OCT dataset. It achieves state-of-the-art performance on retinal disease classification.

Architecture

• Backbone: EfficientNet-B3 (pretrained on ImageNet)
• Head: Dropout(0.3) → Linear(1536→512) → ReLU → Dropout(0.15) → Linear(512→4)
• Input Size: 224×224 RGB images
• Output: 4-class probability distribution

Performance

Metric	Validation	Test
Accuracy	97.8%	99.6%
Macro F1	96.8%	99.6%

Per-Class Performance (Test Set)

Class	Precision	Recall	F1-Score	Support
CNV	98.4%	100.0%	99.2%	242
DME	100.0%	100.0%	100.0%	242
DRUSEN	100.0%	98.4%	99.2%	242
NORMAL	100.0%	100.0%	100.0%	242

Usage

Quick Start

import torch
from PIL import Image
from torchvision import transforms

# Load model
model = torch.hub.load('your-username/oct-classifier', 'oct_classifier', pretrained=True)
model.eval()

# Preprocess image
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Predict
image = Image.open('oct_scan.jpeg').convert('RGB')
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    output = model(input_tensor)
    probabilities = torch.softmax(output, dim=1)
    predicted_class = torch.argmax(probabilities, dim=1).item()

classes = ['CNV', 'DME', 'DRUSEN', 'NORMAL']
print(f"Prediction: {classes[predicted_class]}")
print(f"Confidence: {probabilities[0][predicted_class]:.2%}")

Using with timm

import timm
import torch

# Load the model
model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=4)
model.load_state_dict(torch.load('pytorch_model.bin'))
model.eval()

Training Details

Dataset

• Source: Kermany2018 OCT Dataset
• Total Images: ~84,000 OCT scans
• Train/Val Split: 80/20 stratified
• Test Set: 968 images (242 per class)

Training Configuration

• Optimizer: AdamW (lr=1e-4, weight_decay=0.01)
• Scheduler: Cosine Annealing with 2-epoch warmup
• Batch Size: 32
• Epochs: 20
• Mixed Precision: FP16
• Gradient Clipping: 1.0

Data Augmentation

• Horizontal flip (p=0.5)
• Rotation (±15°)
• Brightness/Contrast adjustment (±0.2)
• Gaussian noise (p=0.3)
• Gaussian blur (p=0.2)

Interpretability

The model includes Grad-CAM visualization support to show which regions of the retinal scan influenced the prediction - critical for clinical trust and validation.

Limitations

• Trained on a single OCT device type; may require fine-tuning for other devices
• Not validated for clinical deployment without proper medical device certification
• Performance may vary on images with different quality or preprocessing

Links

• GitHub: tomalmog/retinal-oct-classifier
• Hugging Face: tomalmog/oct-retinal-classifier

Citation

If you use this model, please cite:

@misc{almog2025oct,
  title={Retinal OCT Disease Classification with EfficientNet-B3},
  author={Almog, Tom},
  year={2025},
  url={https://huggingface.co/tomalmog/oct-retinal-classifier}
}

This model was trained on the Kermany2018 dataset:

@article{kermany2018identifying,
  title={Identifying medical diagnoses and treatable diseases by image-based deep learning},
  author={Kermany, Daniel S and Goldbaum, Michael and Cai, Wenjia and others},
  journal={Cell},
  volume={172},
  number={5},
  pages={1122--1131},
  year={2018},
  publisher={Elsevier}
}

License

MIT License