File size: 3,334 Bytes
33be17a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """Minimal inference example for DoB24/fundus-9model-benchmark.
Loads one of the 9 trained checkpoints, applies the same preprocessing used
during evaluation, and predicts a class for a single fundus image.
Usage:
python inference_example.py path/to/image.jpg --model densenet121
"""
from __future__ import annotations
import argparse, json
from pathlib import Path
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import timm
from huggingface_hub import hf_hub_download
REPO = "DoB24/fundus-9model-benchmark"
CLASSES = [
"Central Serous Chorioretinopathy",
"Diabetic Retinopathy", "Disc Edema", "Glaucoma", "Healthy",
"Macular Scar", "Myopia", "Pterygium",
"Retinal Detachment", "Retinitis Pigmentosa",
]
# Map repo-name to timm model id
TIMM_ID = {
"vgg19": "vgg19", "resnet50": "resnet50", "resnet101": "resnet101",
"densenet121": "densenet121", "inception_v3": "inception_v3",
"swin_b": "swin_base_patch4_window7_224",
}
def load_model(name: str, num_classes: int = 10):
weights = hf_hub_download(REPO, f"weights/{name}_best.pth")
if name in TIMM_ID:
model = timm.create_model(TIMM_ID[name], pretrained=False, num_classes=num_classes)
if name == "inception_v3":
model = timm.create_model("inception_v3", pretrained=False,
num_classes=num_classes, aux_logits=False)
elif name == "clip_openai":
import open_clip
model, _, _ = open_clip.create_model_and_transforms("ViT-B-16", pretrained=None)
model.visual.proj = None
model = torch.nn.Sequential(model.visual, torch.nn.Linear(768, num_classes))
elif name == "dinov2_l":
model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14")
model = torch.nn.Sequential(model, torch.nn.Linear(1024, num_classes))
elif name == "retfound":
model = timm.create_model("vit_large_patch16_224", pretrained=False,
num_classes=num_classes)
else:
raise ValueError(f"Unknown model: {name}")
sd = torch.load(weights, map_location="cpu", weights_only=False)
if isinstance(sd, dict) and "state_dict" in sd:
sd = sd["state_dict"]
model.load_state_dict(sd, strict=False)
model.eval()
return model
def preprocess(img_path: str) -> torch.Tensor:
tf = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
return tf(Image.open(img_path).convert("RGB")).unsqueeze(0)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("image")
ap.add_argument("--model", default="densenet121",
choices=list(TIMM_ID) + ["clip_openai", "dinov2_l", "retfound"])
args = ap.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model(args.model).to(device)
x = preprocess(args.image).to(device)
with torch.no_grad():
probs = F.softmax(model(x), dim=1)[0].cpu().numpy()
order = probs.argsort()[::-1]
print(f"\nTop-3 predictions for {args.image} using {args.model}:")
for i in order[:3]:
print(f" {CLASSES[i]:40s} {probs[i]*100:5.2f} %")
if __name__ == "__main__":
main()
|