retinasense-vit

RetinaSense-ViT v4: Hybrid Vision Transformer for Retinal Disease Classification

RetinaSense-ViT v4 is a hybrid deep learning model that fuses ViT-Base/16 and EfficientNet-B3 feature representations for automated retinal disease screening from color fundus photographs. It classifies images into 5 disease categories with calibrated confidence estimates, uncertainty quantification, and explainable attention maps.

HuggingFace Repository: tanishq74/retinasense-vit

---

Model Details

Architecture Overview

The HybridRetinaModel extracts complementary features from two pretrained backbones and fuses them through concatenation followed by a classification MLP:

Input Image (224x224x3)
    |
    +---> EfficientNet-B3 ---> 1536-dim features
    |                                |
    +---> ViT-Base/16 ---------> 768-dim features
                                     |
                          Concatenation (2304-dim)
                                     |
                              MLP Classifier
                                     |
                           5-class predictions

• EfficientNet-B3 branch: Extracts local texture and fine-grained pathological features (1536-dim pooled features)
• ViT-Base/16 branch: Captures global spatial relationships and long-range dependencies (768-dim CLS token)
• Fusion layer: Concatenated 2304-dim representation passed through a multi-layer perceptron with batch normalization and dropout
• Transfer learning: Initialized from v3 pretrained ViT (best_model.pth, 82.59% val acc) and EfficientNet-B3 (efficientnet_b3.pth, 71.1% val acc)

Disease Classes

---

Training

Dataset

The training dataset was constructed by merging two public retinal imaging datasets:

Balancing was achieved through a combination of undersampling over-represented classes and augmentation-based oversampling of under-represented classes to reach exactly 2,000 images per class.

Data Splits

All splits are patient-aware to prevent data leakage -- no patient appears in more than one split.

Preprocessing Pipeline

The preprocessing pipeline must be replicated exactly for correct inference:

Raw Image -> Crop Black Borders -> Resize 224x224 -> CLAHE (L-channel) -> Circular Mask -> Normalize

Training Configuration

---

Results

5-Fold Stratified Cross-Validation (10,000 images)

Held-Out Test Set (1,486 samples)

Per-Class Test Set Performance

Improvement over v3

---

Usage

Installation

pip install torch torchvision timm opencv-python-headless numpy huggingface_hub

Download Model Weights

from huggingface_hub import hf_hub_download
import shutil, os

os.makedirs("weights", exist_ok=True)
for fname in ["best_model.pth", "temperature.json", "thresholds.json"]:
    path = hf_hub_download(repo_id="tanishq74/retinasense-vit", filename=fname)
    shutil.copy(path, f"weights/{fname}")

Inference Example

import torch
import torch.nn as nn
import timm
import cv2
import json
import numpy as np
from torchvision import transforms

# --- Model Definition ---
class HybridRetinaModel(nn.Module):
    def __init__(self, num_classes=5, drop_rate=0.3):
        super().__init__()
        # EfficientNet-B3 branch (1536-dim)
        self.efficientnet = timm.create_model(
            'efficientnet_b3', pretrained=False, num_classes=0
        )
        # ViT-Base/16 branch (768-dim)
        self.vit = timm.create_model(
            'vit_base_patch16_224', pretrained=False, num_classes=0
        )
        # Fusion MLP: 1536 + 768 = 2304 -> num_classes
        fusion_dim = 2304
        self.classifier = nn.Sequential(
            nn.Linear(fusion_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(drop_rate),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(drop_rate * 0.67),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        eff_features = self.efficientnet(x)   # (B, 1536)
        vit_features = self.vit(x)            # (B, 768)
        fused = torch.cat([eff_features, vit_features], dim=1)  # (B, 2304)
        return self.classifier(fused)

# --- Preprocessing ---
MEAN = [0.4298, 0.2784, 0.1559]
STD = [0.2857, 0.2065, 0.1465]
CLASS_NAMES = ['Normal', 'Diabetes/DR', 'Glaucoma', 'Cataract', 'AMD']

normalize = transforms.Compose([
    transforms.ToPILImage(),
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])

def preprocess(img_path):
    img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
    # Crop black borders
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    mask = gray > 7
    rows, cols = np.any(mask, axis=1), np.any(mask, axis=0)
    if rows.any() and cols.any():
        r0, r1 = np.where(rows)[0][[0, -1]]
        c0, c1 = np.where(cols)[0][[0, -1]]
        img = img[r0:r1+1, c0:c1+1]
    # Resize
    img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_AREA)
    # CLAHE on L-channel
    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    lab[:, :, 0] = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(lab[:, :, 0])
    img = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
    # Circular mask
    h, w = img.shape[:2]
    cmask = np.zeros((h, w), dtype=np.uint8)
    cv2.circle(cmask, (w // 2, h // 2), int(min(h, w) * 0.48), 255, -1)
    img = cv2.bitwise_and(img, img, mask=cmask)
    return normalize(np.clip(img, 0, 255).astype(np.uint8)).unsqueeze(0)

# --- Load Model ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = HybridRetinaModel(num_classes=5).to(device)

checkpoint = torch.load('weights/best_model.pth', map_location=device, weights_only=False)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Load temperature for calibrated probabilities
with open('weights/temperature.json') as f:
    temperature = json.load(f)['temperature']

# Load per-class optimized thresholds (optional)
with open('weights/thresholds.json') as f:
    thresholds = json.load(f)

# --- Run Inference ---
x = preprocess('your_fundus_image.jpg').to(device)

with torch.no_grad():
    logits = model(x)
    probs = torch.softmax(logits / temperature, dim=1).cpu().numpy()[0]

pred_class = CLASS_NAMES[probs.argmax()]
confidence = probs.max()

print(f"Prediction: {pred_class} ({confidence * 100:.1f}%)")
print("\nClass Probabilities:")
for i, name in enumerate(CLASS_NAMES):
    print(f"  {name}: {probs[i] * 100:.1f}%")

MC Dropout Uncertainty Estimation

def mc_dropout_predict(model, x, n_passes=30):
    """Enable dropout at test time for uncertainty estimation."""
    model.train()  # Enable dropout
    predictions = []
    with torch.no_grad():
        for _ in range(n_passes):
            logits = model(x)
            probs = torch.softmax(logits / temperature, dim=1)
            predictions.append(probs.cpu().numpy())
    model.eval()

    predictions = np.array(predictions)  # (n_passes, batch, classes)
    mean_probs = predictions.mean(axis=0)
    epistemic_uncertainty = predictions.var(axis=0).sum(axis=-1)  # Inter-pass variance
    return mean_probs, epistemic_uncertainty

mean_probs, uncertainty = mc_dropout_predict(model, x)
print(f"Prediction: {CLASS_NAMES[mean_probs[0].argmax()]}")
print(f"Epistemic Uncertainty: {uncertainty[0]:.4f}")

---

Features

GradCAM Attention Visualization

Generate class-discriminative heatmaps highlighting the retinal regions most relevant to the model's prediction. Useful for clinical interpretability and verifying that the model attends to pathologically relevant structures.

MC Dropout Uncertainty Estimation

Monte Carlo Dropout with 30 stochastic forward passes provides calibrated epistemic uncertainty estimates. At 90% retention (discarding the most uncertain 10% of predictions), accuracy improves from 80.9% to 86.0%.

Temperature-Calibrated Probabilities

Post-hoc temperature scaling calibrates the model's softmax outputs so that predicted confidence values more closely match empirical accuracy.

Per-Class Optimized Decision Thresholds

Class-specific decision thresholds optimized on the validation set improve macro F1 from 0.813 to 0.822 on the held-out test set.

FAISS-Based Similar Case Retrieval

A FAISS index built from the model's 2304-dimensional fusion embeddings enables fast nearest-neighbor retrieval of visually and diagnostically similar cases from the training set, supporting clinical decision-making through case-based reasoning.

---

Files in Repository

---

Limitations and Ethical Considerations

Technical Limitations

• Normal class has lower F1 (0.69): Early-stage disease presentations overlap visually with normal fundus images, leading to lower precision for the Normal class. Clinically, this means the model errs on the side of flagging images for review rather than missing disease.
• Dataset scope: Trained exclusively on APTOS-2019 and ODIR-5K datasets. Performance may degrade on fundus images from different camera systems, patient demographics, or imaging protocols not represented in the training data.
• Single-label classification: Each image receives one predicted label. Co-morbid conditions (e.g., concurrent DR and glaucoma) are not modeled.
• Cross-validation vs. test gap: The 5-fold CV accuracy (91.1%) is higher than the held-out test accuracy (82.0%), which may reflect distribution differences between augmented training data and real test images.

Intended Use

This model is intended for research and educational purposes only. It may serve as a screening aid or decision-support tool in research settings, but it is not a medical device and has not been validated for clinical deployment.

Out-of-Scope Uses

• Clinical diagnosis without ophthalmologist verification
• Deployment as a standalone screening tool in any healthcare setting
• Use on non-fundus images or imaging modalities other than color fundus photography
• Medico-legal decision-making

---

Citation

@misc{retinasense_v4_2026,
  title={RetinaSense-ViT v4: Hybrid Vision Transformer for Retinal Disease Classification},
  author={Tanishq Tamarkar},
  year={2026},
  url={https://huggingface.co/tanishq74/retinasense-vit},
  note={ViT-Base/16 + EfficientNet-B3 hybrid fusion model, 5-class retinal disease classification}
}

---

License

MIT License

---

Disclaimer

This is a research prototype for AI-assisted retinal screening. It is NOT a certified medical device and should NOT be used for clinical decision-making without independent verification by a qualified ophthalmologist. All predictions are probabilistic estimates and may be incorrect. The authors assume no liability for any clinical decisions made based on this model's outputs.