""" 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 import os import sys # Add parent directory to path to find 'utils' sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 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."