models/predict.py

"""
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."