| """DR grader: loads the trained EfficientNet-B4 + Swin (CORN) model and grades |
| a fundus image into one of five diabetic-retinopathy severity classes. |
| |
| The model architecture and weights are reused from the original DR Grading |
| project; the weights are downloaded from the public HuggingFace Space on first |
| use. Inference uses the training-matched preprocessing (crop + CLAHE). |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| import timm |
| from torchvision import transforms |
| from huggingface_hub import hf_hub_download |
|
|
| from src.config import ( |
| DR_WEIGHTS_REPO, DR_WEIGHTS_REPO_TYPE, DR_WEIGHTS_FILE, |
| EFFICIENTNET_NAME, SWIN_NAME, NUM_CLASSES, FUSION_HIDDEN_DIM, FUSION_DROPOUT, |
| IMAGE_SIZE, NORM_MEAN, NORM_STD, CLASS_NAMES, |
| ) |
| from src.preprocessing import preprocess_fundus |
|
|
|
|
| class ParallelEfficientNetSwinCORN(nn.Module): |
| """EfficientNet-B4 + Swin Transformer parallel backbones with a CORN ordinal head. |
| |
| Must match the trained architecture exactly (layer names included) so the |
| saved weights load with strict=True. |
| """ |
|
|
| def __init__(self, efficientnet_name, swin_name, num_classes, |
| fusion_hidden_dim=1024, fusion_dropout=0.3, pretrained=False): |
| super().__init__() |
| self.efficientnet = timm.create_model(efficientnet_name, pretrained=pretrained, num_classes=0) |
| self.swin = timm.create_model(swin_name, pretrained=pretrained, num_classes=0) |
| eff_dim = getattr(self.efficientnet, "num_features", None) |
| swin_dim = getattr(self.swin, "num_features", None) |
| if eff_dim is None or swin_dim is None: |
| raise ValueError("Could not determine backbone feature dimensions.") |
| fusion_dim = eff_dim + swin_dim |
| self.fusion = nn.Sequential( |
| nn.LayerNorm(fusion_dim), |
| nn.Linear(fusion_dim, fusion_hidden_dim), |
| nn.GELU(), |
| nn.Dropout(fusion_dropout), |
| nn.Linear(fusion_hidden_dim, num_classes - 1), |
| ) |
|
|
| @staticmethod |
| def _to_vector(feat): |
| if feat.ndim == 2: |
| return feat |
| if feat.ndim == 3: |
| return feat.mean(dim=1) |
| if feat.ndim == 4: |
| if feat.shape[1] > feat.shape[-1] and feat.shape[1] > feat.shape[-2]: |
| return feat.mean(dim=(2, 3)) |
| return feat.mean(dim=(1, 2)) |
| raise ValueError(f"Unexpected feature shape: {feat.shape}") |
|
|
| def forward(self, x): |
| eff = self._to_vector(self.efficientnet(x)) |
| swin = self._to_vector(self.swin(x)) |
| return self.fusion(torch.cat([eff, swin], dim=1)) |
|
|
|
|
| def _clean_state_dict(state_dict): |
| """Strip any 'module.' prefixes left by DataParallel during training.""" |
| return {(k[len("module."):] if k.startswith("module.") else k): v |
| for k, v in state_dict.items()} |
|
|
|
|
| _model = None |
| _device = None |
| _transform = transforms.Compose([ |
| transforms.ToTensor(), |
| transforms.Normalize(NORM_MEAN, NORM_STD), |
| ]) |
|
|
|
|
| def load_grader(): |
| """Load the model + weights once (cached for repeated calls).""" |
| global _model, _device |
| if _model is not None: |
| return _model, _device |
| _device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"Downloading DR weights from {DR_WEIGHTS_REPO} (first run only) ...") |
| weights_path = hf_hub_download( |
| repo_id=DR_WEIGHTS_REPO, filename=DR_WEIGHTS_FILE, repo_type=DR_WEIGHTS_REPO_TYPE, |
| ) |
| model = ParallelEfficientNetSwinCORN( |
| efficientnet_name=EFFICIENTNET_NAME, swin_name=SWIN_NAME, |
| num_classes=NUM_CLASSES, fusion_hidden_dim=FUSION_HIDDEN_DIM, |
| fusion_dropout=FUSION_DROPOUT, pretrained=False, |
| ) |
| state_dict = _clean_state_dict(torch.load(weights_path, map_location=_device)) |
| model.load_state_dict(state_dict, strict=True) |
| model.to(_device).eval() |
| _model = model |
| return _model, _device |
|
|
|
|
| def grade_fundus(pil_image): |
| """Grade one fundus image. |
| |
| Returns the grade (0-4), its label, and the cumulative ordinal |
| probabilities P(y > k). CORN decoding: the K-1 logits are conditional |
| probabilities, so the cumulative P(y > k) is the cumulative product of |
| their sigmoids, and the predicted grade is the count of cumulative |
| probabilities above 0.5 (equivalent to coral_pytorch.corn_label_from_logits). |
| """ |
| model, device = load_grader() |
| processed = preprocess_fundus(pil_image) |
| x = _transform(processed).unsqueeze(0).to(device) |
| with torch.inference_mode(): |
| logits = model(x) |
| |
| |
| |
| |
| probs = torch.cumprod(torch.sigmoid(logits), dim=1) |
| grade = int((probs > 0.5).sum(dim=1).item()) |
| return { |
| "grade": grade, |
| "label": CLASS_NAMES[grade], |
| "ordinal_probabilities": [round(p, 4) for p in probs.squeeze(0).cpu().tolist()], |
| "processed_image": processed, |
| } |
|
|
|
|
| if __name__ == "__main__": |
| |
| import sys |
| from PIL import Image |
|
|
| if len(sys.argv) < 2: |
| print("Usage: python -m src.dr_grader <path-to-fundus-image>") |
| raise SystemExit(1) |
|
|
| result = grade_fundus(Image.open(sys.argv[1])) |
| print(f"\nGrade {result['grade']} - {result['label']}") |
| print(f"Ordinal probabilities P(>0..>3): {result['ordinal_probabilities']}") |
|
|
