Spaces:
Sleeping
Sleeping
| import os, sys, torch, joblib, numpy as np, gradio as gr, spaces | |
| from PIL import Image | |
| from torchvision import transforms | |
| from huggingface_hub import snapshot_download | |
| from util.pos_embed import interpolate_pos_embed | |
| from timm.models.layers import trunc_normal_ | |
| import models_vit | |
| print("Downloading models...") | |
| snapshot_download( | |
| repo_id="janvib/retinal-weights", | |
| local_dir="/tmp/models", | |
| repo_type="model", | |
| ) | |
| print("Downloaded.") | |
| pth_files = [f for f in os.listdir("/tmp/models") if f.endswith(".pth")] | |
| weights_path = f"/tmp/models/{pth_files[0]}" | |
| model = models_vit.RETFound_mae(num_classes=2, drop_path_rate=0.2, global_pool=True) | |
| checkpoint = torch.load(weights_path, map_location="cpu", weights_only=False) | |
| checkpoint_model = checkpoint["model"] | |
| state_dict = model.state_dict() | |
| for k in ["head.weight", "head.bias"]: | |
| if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape: | |
| del checkpoint_model[k] | |
| interpolate_pos_embed(model, checkpoint_model) | |
| model.load_state_dict(checkpoint_model, strict=False) | |
| model = model.eval() | |
| ridge = joblib.load("/tmp/models/retinal_age_ridge.pkl") | |
| glaucoma_clf = joblib.load("/tmp/models/glaucoma_classifier.pkl") | |
| dr_clf = joblib.load("/tmp/models/dr_classifier_v2.pkl") | |
| dr_scaler = joblib.load("/tmp/models/dr_scaler.pkl") | |
| print("All models loaded") | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def predict(image, age): | |
| img_tensor = transform(image).unsqueeze(0).cuda() | |
| model.cuda() | |
| with torch.no_grad(): | |
| feats = model.forward_features(img_tensor)[:, 0, :].cpu().numpy() | |
| feats_scaled = dr_scaler.transform(feats) | |
| g_prob = glaucoma_clf.predict_proba(feats)[0][1] | |
| dr_prob = dr_clf.predict_proba(feats_scaled)[0][1] | |
| predicted_age = float(ridge.predict(feats)[0]) | |
| rag = predicted_age - float(age) if age else None | |
| if g_prob >= 0.7: | |
| g_flag = "HIGH - Referable Glaucoma Suspected" | |
| g_rec = "Recommend full glaucoma workup: IOP, visual fields, OCT RNFL." | |
| elif g_prob >= 0.5: | |
| g_flag = "MODERATE - Glaucoma Suspect" | |
| g_rec = "Consider glaucoma suspect workup." | |
| else: | |
| g_flag = "LOW - No Referable Glaucoma" | |
| g_rec = "Routine follow-up." | |
| if dr_prob >= 0.7: | |
| dr_flag = "HIGH - Referable Diabetic Retinopathy" | |
| dr_rec = "Recommend urgent retinal specialist referral." | |
| elif dr_prob >= 0.5: | |
| dr_flag = "MODERATE - Possible Referable DR" | |
| dr_rec = "Consider retinal specialist referral." | |
| else: | |
| dr_flag = "LOW - No Referable DR" | |
| dr_rec = "Routine follow-up if diabetic." | |
| age_line = "Predicted Retinal Age: " + str(round(predicted_age, 1)) + " yrs" | |
| if rag is not None: | |
| age_line = age_line + "\nRetinal Age Gap: " + str(round(rag, 1)) + " yrs" | |
| result = "\nGLAUCOMA SCREENING\n" | |
| result += g_flag + "\n" | |
| result += "Probability: " + str(round(g_prob * 100, 1)) + "%\n" | |
| result += "Action: " + g_rec + "\n" | |
| result += "\nDIABETIC RETINOPATHY SCREENING\n" | |
| result += dr_flag + "\n" | |
| result += "Probability: " + str(round(dr_prob * 100, 1)) + "%\n" | |
| result += "Action: " + dr_rec + "\n" | |
| result += "\nRETINAL AGE\n" | |
| result += age_line + "\n" | |
| result += "\nResearch use only. Not FDA-cleared. Clinical correlation required." | |
| return result | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Image(type="pil", label="Fundus Photograph"), | |
| gr.Number(label="Patient Age (optional)", value=60), | |
| ], | |
| outputs=gr.Textbox(label="Screening Results", lines=22), | |
| title="RETFound - Retinal Screening Tool", | |
| description="Upload a de-identified color fundus photograph to screen for glaucoma and diabetic retinopathy. Research use only. Not FDA-cleared.", | |
| flagging_mode="never", | |
| ) | |
| demo.launch() | |