Downloads · 30 days
4
8% of all-time downloads
pooja-arumugam/cataract-vit
cataract-vit is a image classification model from pooja-arumugam. Use it when you need a label for an image. The card lists the license as mit.
Hybrid ResNet50 + Vision Transformer for classifying cataract severity into 4 classes.
Downloads · 30 days
4
8% of all-time downloads
All-time downloads
50
Public
Repo size
132 MB
Likes
0
Public
Click a slice to open those files.
.pth132 MB · 100%
From the Hugging Face model README
Hybrid ResNet50 + Vision Transformer for classifying cataract severity into 4 classes.
| Label | Description |
|---|---|
| Normal | No cataract |
| Mild | Early stage |
| Moderate | Medium stage |
| Severe | Mature cataract |
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()])