Initial commit

README.md

@@ -1,4 +1,3 @@
-<<<<<<< HEAD
 ---
 license: apache-2.0
 language: en
@@ -13,9 +12,9 @@ tags:
 datasets:
 - aptos2019-blindness-detection
 widget:
-- src: model/gradcam_visualizations/gradcam_sample_003.png
+- src: gradcam_visualizations/gradcam_sample_003.png
   example_title: No DR Example
-- src: model/gradcam_visualizations/gradcam_sample_007.png
+- src: gradcam_visualizations/gradcam_sample_007.png
   example_title: Severe DR Example
 ---
 
@@ -41,9 +40,184 @@ The model can be easily loaded from Hugging Face Hub for inference.
 
 ```bash
 # Install required libraries
-pip install torch torchvision timm albumentations huggingface-hub numpy<2.0 pillow opencv-python
-=======
----
-license: apache-2.0
----
->>>>>>> 96df5800dd01e6edcea4c1417179b056d6e3e22d
+pip install torch torchvision timm albumentations huggingface-hub numpy pillow opencv-python
+```
+
+```python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+from PIL import Image
+import numpy as np
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from huggingface_hub import hf_hub_download
+
+# Define the model architecture
+class MultiTaskDRModel(nn.Module):
+    def __init__(self, model_name='efficientnet_b3', num_classes=5,
+                 num_lesion_types=5, num_regions=5, pretrained=False):
+        super(MultiTaskDRModel, self).__init__()
+        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
+        self.feature_dim = self.backbone.num_features
+        
+        self.attention = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
+            nn.Linear(self.feature_dim, self.feature_dim // 8), nn.ReLU(inplace=True),
+            nn.Linear(self.feature_dim // 8, self.feature_dim), nn.Sigmoid()
+        )
+        
+        self.feature_norm = nn.BatchNorm1d(self.feature_dim)
+        self.dropout = nn.Dropout(0.4)
+        
+        self.severity_classifier = nn.Sequential(
+            nn.Linear(self.feature_dim, self.feature_dim // 2), nn.ReLU(inplace=True),
+            nn.Dropout(0.2), nn.Linear(self.feature_dim // 2, num_classes)
+        )
+        
+        self.lesion_detector = nn.Sequential(
+            nn.Linear(self.feature_dim, self.feature_dim // 4), nn.ReLU(inplace=True),
+            nn.Dropout(0.2), nn.Linear(self.feature_dim // 4, num_lesion_types)
+        )
+        
+        self.region_predictor = nn.Sequential(
+            nn.Linear(self.feature_dim, self.feature_dim // 4), nn.ReLU(inplace=True),
+            nn.Dropout(0.2), nn.Linear(self.feature_dim // 4, num_regions)
+        )
+    
+    def forward(self, x):
+        features = self.backbone.forward_features(x)
+        pooled_features = F.adaptive_avg_pool2d(features, 1).flatten(1)
+        attention_weights = self.attention(pooled_features.unsqueeze(-1).unsqueeze(-1))
+        features = pooled_features * attention_weights
+        features = self.feature_norm(features)
+        features = self.dropout(features)
+        
+        severity_logits = self.severity_classifier(features)
+        lesion_logits = self.lesion_detector(features)
+        region_logits = self.region_predictor(features)
+        
+        return {
+            'severity': severity_logits,
+            'lesions': lesion_logits,
+            'regions': region_logits,
+            'features': features
+        }
+
+# Load the model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = MultiTaskDRModel()
+
+# Download and load the checkpoint
+model_path = hf_hub_download(
+    repo_id="dheeren-tejani/DiabeticRetinpathyClassifier",
+    filename="best_model_v2.pth"
+)
+checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+model.load_state_dict(checkpoint['model_state_dict'])
+model.to(device)
+model.eval()
+
+print("Model loaded successfully!")
+
+# Preprocessing function
+def preprocess_image(image_path):
+    transforms = A.Compose([
+        A.Resize(512, 512),
+        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ToTensorV2(),
+    ])
+    image = np.array(Image.open(image_path).convert("RGB"))
+    image_tensor = transforms(image=image)['image'].unsqueeze(0)
+    return image_tensor
+
+# Example inference
+def predict_dr_severity(image_path):
+    image_tensor = preprocess_image(image_path).to(device)
+    
+    with torch.no_grad():
+        outputs = model(image_tensor)
+    
+    # Get severity prediction
+    severity_probs = torch.softmax(outputs['severity'], dim=1)
+    predicted_class = torch.argmax(severity_probs, dim=1).item()
+    confidence = severity_probs[0, predicted_class].item()
+    
+    severity_labels = {
+        0: "No DR",
+        1: "Mild DR", 
+        2: "Moderate DR",
+        3: "Severe DR",
+        4: "Proliferative DR"
+    }
+    
+    return {
+        'predicted_severity': severity_labels[predicted_class],
+        'confidence': confidence,
+        'all_probabilities': severity_probs[0].cpu().numpy()
+    }
+
+# Example usage
+# result = predict_dr_severity("path/to/your/fundus_image.jpg")
+# print(f"Predicted: {result['predicted_severity']} (Confidence: {result['confidence']:.3f})")
+```
+
+## Training Details
+
+### V2 Improvements
+This model (V2) was specifically designed to address the shortcomings of a baseline model (V1) that struggled with severe-stage DR detection:
+
+- **Higher Resolution:** Increased from 224×224 to 512×512 to capture finer pathological details
+- **Class Balancing:** Implemented WeightedRandomSampler to oversample rare minority classes (Severe and Proliferative DR)
+- **Focal Loss:** Replaced standard Cross-Entropy with Focal Loss (γ=2.0) to focus on hard-to-classify examples
+- **Focused Training:** Set auxiliary task weights to zero, dedicating full model capacity to severity classification
+
+### Hyperparameters
+- **Optimizer:** AdamW
+- **Learning Rate:** 1e-4
+- **Scheduler:** CosineAnnealingWarmRestarts (T_MAX=10)
+- **Batch Size:** 16
+- **Epochs:** 17 (Early stopping)
+- **Image Size:** 512×512
+
+## Performance
+
+The model was evaluated on a held-out validation set of 735 images:
+
+| Metric | Score |
+|--------|-------|
+| **Quadratic Weighted Kappa (QWK)** | **0.796** |
+| Accuracy | 65.0% |
+| F1-Score (Weighted) | 66.3% |
+| F1-Score (Macro) | 53.5% |
+
+### Key Achievement
+The V2 model achieved a **+3.5% improvement in QWK** over the V1 baseline (0.761), indicating it makes "smarter" errors that are more aligned with clinical judgment, despite lower overall accuracy. This trade-off prioritizes clinically relevant performance over naive accuracy.
+
+## Limitations
+
+⚠️ **Important Disclaimers:**
+- This model was trained on a single public dataset and may not generalize to different clinical settings, camera types, or patient demographics
+- The dataset may contain inherent demographic biases
+- **This is NOT a medical device** and should not be used for actual clinical diagnosis
+- Always consult qualified healthcare professionals for medical decisions
+
+## Citation
+
+If you use this model in your research, please cite:
+
+```bibtex
+@misc{dheerentejani2025dr,
+  author = {Dheeren Tejani},
+  title = {Diabetic Retinopathy Grading Model V2},
+  year = {2025},
+  publisher = {Hugging Face},
+  journal = {Hugging Face Model Hub},
+  howpublished = {\url{https://huggingface.co/dheeren-tejani/DiabeticRetinpathyClassifier}},
+}
+```
+
+## License
+
+This model is released under the Apache 2.0 License.