Diabetic-Retinopathy-Detection

MobileNetV2 + PiT Transformer for Diabetic Retinopathy Detection

This repository contains a deep learning model combining MobileNetV2 and PiT (Pooling-based Vision Transformer) for the early detection of diabetic retinopathy from retinal fundus images.
The model is trained using the Lambda optimizer and achieves a high test accuracy of 99.01%, making it suitable for clinical and real-world deployment.

---

Model Description

• Task: Early detection of Diabetic Retinopathy (DR)
• Architecture: MobileNetV2 backbone + PiT (Pooling-based Vision Transformer)
• Optimizer: Lambda optimizer
• Dataset: Retinal fundus images (preprocessed)
• Performance: 99.01% test accuracy
• Framework: PyTorch
• Input: RGB retinal image (size depends on training setup, e.g., 224×224)
• Output: DR stage classification (e.g., No DR, Mild, Moderate, Severe, Proliferative)

---

🔗 Links

• GitHub Repository: GitHub Repo

---

🩺 Use Case

This model is designed to assist in automated screening of diabetic retinopathy, enabling:

• Early diagnosis and prevention of vision loss
• Fast and accurate classification of DR stages
• Integration into telemedicine platforms and clinical workflows

---

Usage

You can load the model and run inference on a retinal fundus image as shown below:

import torch
import torch.nn as nn
import timm
from torchvision import models, transforms
from PIL import Image
from huggingface_hub import hf_hub_download

# ---------- DEVICE ----------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------- CLASS LABELS ----------
class_names = ['No_DR', 'Mild', 'Moderate', 'Severe', 'Proliferate_DR']

# ---------- MODEL ARCHITECTURE ----------
Model_P = "pit_b_224"
Model_Mobile = models.mobilenet_v2(pretrained=True)
num_features_mobile = Model_Mobile.classifier[1].in_features
Model_Mobile.classifier = nn.Identity()

Model_pit = timm.create_model(Model_P, pretrained=False, num_classes=5)
num_features_pit = Model_pit.head.in_features
Model_pit.head = nn.Identity()

class CombinedModel(nn.Module):
    def __init__(self, Model_Mobile, Model_pit, num_classes):
        super(CombinedModel, self).__init__()
        self.Model_Mobile = Model_Mobile
        self.Model_pit = Model_pit
        self.fc = nn.Linear(num_features_mobile + num_features_pit, num_classes)

    def forward(self, x):
        mobile_features = self.Model_Mobile(x)
        pit_features = self.Model_pit(x)

        if len(mobile_features.shape) > 2:
            mobile_features = mobile_features.flatten(1)
        if len(pit_features.shape) > 2:
            pit_features = pit_features.flatten(1)

        combined = torch.cat((mobile_features, pit_features), dim=1)
        return self.fc(combined)

# ---------- LOAD MODEL FROM HUGGING FACE ----------
def load_model():
    repo_id = "Tanishrajput/Diabetic-Retinopathy-Detection"
    filename = "Mobile_Pit_last_checkpoint.pt"
    token = os.environ.get("HUGGINGFACE_TOKEN")  # Hugging Face token from environment

    model = CombinedModel(Model_Mobile, Model_pit, num_classes=5).to(device)
    
    checkpoint_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        token=token
    )
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model

# ---------- PREDICT FUNCTION ----------
def predict_retinal_image(image_path, model):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
    image = Image.open(image_path).convert("RGB")
    input_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(input_tensor)
        predicted_idx = torch.argmax(output, dim=1).item()
    
    return class_names[predicted_idx]

# ---------- USAGE ----------
if __name__ == "__main__":
    model = load_model()
    image_path = "sample_retina.jpg"  # Replace with your image path
    prediction = predict_retinal_image(image_path, model)
    print("Predicted DR Stage:", prediction)

---

License

This model is released under the MIT License.