Spaces:
Sleeping
Sleeping
File size: 8,748 Bytes
5bd1dc2 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | """
Prediction Functions
Handles predictions for Fundus, OCT, and Multimodal inputs
"""
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import numpy as np
import base64
from io import BytesIO
from utils.vcdr import compute_robust_vcdr
from utils.gradcam import generate_gradcam_image
from sklearn.preprocessing import StandardScaler
# Image preprocessing
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
fundus_transform = transforms.Compose([
transforms.Resize((288, 288)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
])
oct_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
])
multimodal_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
])
def image_to_base64(image_path):
"""Convert image to base64 for display"""
with open(image_path, 'rb') as f:
img_data = f.read()
return base64.b64encode(img_data).decode('utf-8')
def predict_fundus(image_path, model_loader):
"""Predict on fundus image"""
device = model_loader.device
model = model_loader.fundus_model
# Load and preprocess image
img = Image.open(image_path).convert('RGB')
img_tensor = fundus_transform(img).unsqueeze(0).to(device)
# Compute vCDR
vcdr = compute_robust_vcdr(image_path)
# Predict
with torch.no_grad():
outputs = model(img_tensor)
probs = F.softmax(outputs, dim=1)
pred_idx = torch.argmax(probs, dim=1).item()
confidence = probs[0][pred_idx].item()
# Generate Grad-CAM
gradcam_b64 = generate_gradcam_image(model, img_tensor, pred_idx, device)
# Prepare result
result = {
'prediction': model_loader.fundus_classes[pred_idx],
'confidence': float(confidence),
'vcdr': float(vcdr),
'probabilities': {
model_loader.fundus_classes[i]: float(probs[0][i])
for i in range(len(model_loader.fundus_classes))
},
'original_image': image_to_base64(image_path),
'gradcam_image': gradcam_b64,
'interpretation': get_fundus_interpretation(pred_idx, confidence, vcdr)
}
return result
def predict_oct(image_path, model_loader):
"""Predict on OCT image"""
device = model_loader.device
model = model_loader.oct_model
# Load and preprocess image
img = Image.open(image_path).convert('RGB')
img_tensor = oct_transform(img).unsqueeze(0).to(device)
# Predict
with torch.no_grad():
outputs = model(img_tensor)
probs = F.softmax(outputs, dim=1)
pred_idx = torch.argmax(probs, dim=1).item()
confidence = probs[0][pred_idx].item()
# Get top 3 predictions
top3_probs, top3_indices = torch.topk(probs[0], k=min(3, len(model_loader.oct_classes)))
top3 = [
{
'class': model_loader.oct_classes[idx.item()],
'probability': float(prob.item())
}
for prob, idx in zip(top3_probs, top3_indices)
]
# Generate Grad-CAM
gradcam_b64 = generate_gradcam_image(model, img_tensor, pred_idx, device)
# Prepare result
result = {
'prediction': model_loader.oct_classes[pred_idx],
'confidence': float(confidence),
'probabilities': {
model_loader.oct_classes[i]: float(probs[0][i])
for i in range(len(model_loader.oct_classes))
},
'top3': top3,
'original_image': image_to_base64(image_path),
'gradcam_image': gradcam_b64,
'interpretation': get_oct_interpretation(pred_idx, confidence, model_loader.oct_classes)
}
return result
def predict_multimodal(fundus_path, oct_path, model_loader):
"""Predict using multimodal fusion"""
device = model_loader.device
model = model_loader.fusion_model
# Load and preprocess images
fundus_img = Image.open(fundus_path).convert('RGB')
oct_img = Image.open(oct_path).convert('RGB')
fundus_tensor = multimodal_transform(fundus_img).unsqueeze(0).to(device)
oct_tensor = multimodal_transform(oct_img).unsqueeze(0).to(device)
# Compute clinical features
vcdr = compute_robust_vcdr(fundus_path)
rnfl = 90.0 # Default RNFL value (would be computed from OCT layers in production)
# Scale clinical features
scaler = StandardScaler()
scaler.mean_ = np.array([0.5, 90.0])
scaler.scale_ = np.array([0.2, 20.0])
clinical_feats = np.array([[vcdr, rnfl]], dtype=np.float32)
clinical_feats = scaler.transform(clinical_feats).flatten()
clinical_tensor = torch.tensor(clinical_feats, dtype=torch.float32).unsqueeze(0).to(device)
# Predict
with torch.no_grad():
outputs = model(fundus_tensor, oct_tensor, clinical_tensor)
probs = F.softmax(outputs, dim=1)
pred_idx = torch.argmax(probs, dim=1).item()
confidence = probs[0][pred_idx].item()
# Generate Grad-CAMs for both modalities
fundus_gradcam = generate_gradcam_image(model.fundus_model, fundus_tensor, pred_idx, device)
oct_gradcam = generate_gradcam_image(model.oct_model, oct_tensor, pred_idx, device)
# Prepare result
result = {
'prediction': model_loader.fusion_classes[pred_idx],
'confidence': float(confidence),
'vcdr': float(vcdr),
'rnfl': float(rnfl),
'probabilities': {
model_loader.fusion_classes[i]: float(probs[0][i])
for i in range(len(model_loader.fusion_classes))
},
'fundus_image': image_to_base64(fundus_path),
'oct_image': image_to_base64(oct_path),
'fundus_gradcam': fundus_gradcam,
'oct_gradcam': oct_gradcam,
'interpretation': get_multimodal_interpretation(pred_idx, confidence, vcdr, rnfl)
}
return result
def get_fundus_interpretation(pred_idx, confidence, vcdr):
"""Generate clinical interpretation for fundus prediction"""
if pred_idx == 1: # GON+
severity = "High" if vcdr > 0.7 else "Moderate" if vcdr > 0.6 else "Mild"
return f"Glaucoma detected with {severity.lower()} severity. vCDR of {vcdr:.2f} indicates {'significant' if vcdr > 0.7 else 'moderate'} optic nerve damage. Recommend immediate ophthalmologist consultation."
else:
return f"No glaucoma detected. vCDR of {vcdr:.2f} is within normal range. Continue regular eye examinations."
def get_oct_interpretation(pred_idx, confidence, class_names):
"""Generate clinical interpretation for OCT prediction"""
disease = class_names[pred_idx]
interpretations = {
'AMD': "Age-related Macular Degeneration detected. This condition affects central vision. Recommend anti-VEGF therapy consultation.",
'DME': "Diabetic Macular Edema detected. Blood sugar control and anti-VEGF treatment may be needed.",
'ERM': "Epiretinal Membrane detected. Monitor for vision changes. Surgery may be considered if vision deteriorates.",
'NO': "No significant retinal pathology detected. Continue routine eye examinations.",
'RAO': "Retinal Artery Occlusion detected. This is a medical emergency requiring immediate treatment.",
'RVO': "Retinal Vein Occlusion detected. Urgent ophthalmologist consultation recommended.",
'VID': "Vitreous Hemorrhage detected. Requires immediate evaluation to determine underlying cause."
}
return interpretations.get(disease, "Consult with ophthalmologist for detailed evaluation.")
def get_multimodal_interpretation(pred_idx, confidence, vcdr, rnfl):
"""Generate clinical interpretation for multimodal prediction"""
if pred_idx == 1: # GLAUCOMA
vcdr_status = "elevated" if vcdr > 0.6 else "borderline"
rnfl_status = "thinning" if rnfl < 85 else "borderline"
return f"Multimodal analysis indicates glaucoma. vCDR is {vcdr_status} at {vcdr:.2f}, and RNFL shows {rnfl_status} at {rnfl:.1f}μm. Combined evidence from fundus and OCT imaging supports this diagnosis. Immediate ophthalmologist consultation strongly recommended for treatment planning."
else:
return f"Multimodal analysis shows no signs of glaucoma. vCDR of {vcdr:.2f} and RNFL of {rnfl:.1f}μm are within normal limits. Continue regular monitoring."
|