pooja-arumugam/cataract-vit

CataractViT — Cataract Severity Classifier

Hybrid ResNet50 + Vision Transformer for classifying cataract severity into 4 classes.

Classes

Label	Description
Normal	No cataract
Mild	Early stage
Moderate	Medium stage
Severe	Mature cataract

Validation Accuracy: 98.57%

Usage

import torch
from PIL import Image
import torchvision.transforms as T
from model import CataractViT

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
CLASS_NAMES = ['mild', 'moderate', 'normal', 'severe']

ckpt  = torch.load('cataract_best_v2.pth', map_location=DEVICE)
model = CataractViT(num_classes=4).to(DEVICE)
model.load_state_dict(ckpt['model'])
model.eval()

tf = T.Compose([
    T.Resize((224,224)), T.ToTensor(),
    T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])
img    = Image.open('eye.jpg').convert('RGB')
tensor = tf(img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
    probs = torch.softmax(model(tensor), dim=1)[0]
print(CLASS_NAMES[probs.argmax()])