inference.py

import json
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
from pathlib import Path

# ─── Architecture (must match training notebook exactly) ───────────────────────

class ChannelAttention(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(channels, channels // reduction, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels // reduction, channels, 1, bias=False),
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        return self.sigmoid(self.fc(self.avg_pool(x)) + self.fc(self.max_pool(x)))


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))


class CBAM(nn.Module):
    def __init__(self, channels, reduction=16, kernel_size=7):
        super().__init__()
        self.channel_attention = ChannelAttention(channels, reduction)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        x = x * self.channel_attention(x)
        x = x * self.spatial_attention(x)
        return x


class DenseNet121ForBinaryClassification(nn.Module):
    def __init__(self, num_classes=1, pretrained=False, dropout=0.5):
        super().__init__()
        self.densenet = models.densenet121(pretrained=pretrained)
        orig_conv = self.densenet.features.conv0
        self.densenet.features.conv0 = nn.Conv2d(
            1, orig_conv.out_channels,
            kernel_size=orig_conv.kernel_size,
            stride=orig_conv.stride,
            padding=orig_conv.padding,
            bias=False
        )
        self.cbam = CBAM(channels=1024, reduction=16, kernel_size=7)
        num_features = self.densenet.classifier.in_features
        self.densenet.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(num_features, num_classes)
        )

    def forward(self, x):
        features = self.densenet.features(x)
        features = self.cbam(features)
        out = F.adaptive_avg_pool2d(features, (1, 1))
        out = out.view(out.size(0), -1)
        return self.densenet.classifier(out)


# ─── GradCAM++ ─────────────────────────────────────────────────────────────────

class GradCAMPlusPlus:
    def __init__(self, model, target_layer_name="cbam"):
        self.model = model
        self.gradients = None
        self.activations = None
        self.hooks = []
        self._register_hooks(target_layer_name)

    def _register_hooks(self, target_layer_name):
        def forward_hook(module, input, output):
            self.activations = output.detach()

        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0].detach()

        for name, module in self.model.named_modules():
            if name == target_layer_name:
                self.hooks.append(module.register_forward_hook(forward_hook))
                self.hooks.append(module.register_backward_hook(backward_hook))
                return
        raise ValueError(f"Layer '{target_layer_name}' not found in model")

    def generate_cam(self, input_tensor, class_idx=None):
        self.model.eval()

        for param in self.model.parameters():
            param.requires_grad = True

        input_tensor = input_tensor.clone().detach().requires_grad_(True)
        output = self.model(input_tensor)

        if class_idx is None:
            class_idx = int(torch.sigmoid(output).round().item())

        target_score = output[0, 0] if class_idx == 1 else -output[0, 0]
        self.model.zero_grad()
        target_score.backward(retain_graph=True)

        grads = self.gradients    # [1, C, H, W]
        acts  = self.activations  # [1, C, H, W]

        B, C, H, W = grads.shape

        # ── GradCAM++ alpha computation (all ops stay in [B, C, H, W]) ──
        grads_sq  = grads.pow(2)                          # [1, C, H, W]
        grads_cub = grads.pow(3)                          # [1, C, H, W]

        # sum over spatial dims H,W → [1, C, 1, 1]  then broadcast back
        spatial_sum = (acts * grads_cub).sum(dim=[2, 3], keepdim=True)  # [1, C, 1, 1]

        alpha_denom = 2.0 * grads_sq + spatial_sum        # [1, C, H, W]
        alpha_denom = torch.where(
            alpha_denom != 0.0,
            alpha_denom,
            torch.ones_like(alpha_denom)
        )
        alpha = grads_sq / (alpha_denom + 1e-7)           # [1, C, H, W]

        # weights: alpha * ReLU(grads), summed over H,W → [1, C, 1, 1]
        weights = (alpha * torch.relu(grads)).sum(dim=[2, 3], keepdim=True)  # [1, C, 1, 1]

        # CAM: weighted sum of activations → [1, 1, H, W]
        cam = (weights * acts).sum(dim=1, keepdim=True)   # [1, 1, H, W]
        cam = torch.relu(cam)

        cam = cam.squeeze().cpu().detach().numpy()         # [H, W]
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        return cam

    def cleanup(self):
        for hook in self.hooks:
            hook.remove()



# ─── Preprocessing ─────────────────────────────────────────────────────────────

def apply_jet_colormap(gray_img):
    """Apply jet colormap manually without matplotlib (0=blue, 1=red)."""
    gray_img = np.clip(gray_img, 0, 1)
    r = np.clip(1.5 - np.abs(gray_img * 2 - 3), 0, 1)
    g = np.clip(1.5 - np.abs(gray_img * 2 - 2), 0, 1)
    b = np.clip(1.5 - np.abs(gray_img * 2 - 1), 0, 1)
    return (np.stack([r, g, b], axis=-1) * 255).astype(np.uint8)


def apply_clahe_cv2(image, clip_limit=2.0, tile_grid_size=(8, 8)):
    """Apply CLAHE using OpenCV (replaces albumentations)."""
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
    return clahe.apply(image)


def preprocess_with_cv2(image, image_size=512):
    """Preprocessing pipeline using pure OpenCV (replaces albumentations)."""
    # CLAHE
    image = apply_clahe_cv2(image, clip_limit=2.0, tile_grid_size=(8, 8))
    
    # Center crop 350x350
    h, w = image.shape[:2]
    start_y = (h - 350) // 2
    start_x = (w - 350) // 2
    image = image[start_y:start_y+350, start_x:start_x+350]
    
    # Resize to target size
    image = cv2.resize(image, (image_size, image_size), interpolation=cv2.INTER_LINEAR)
    
    return image


def preprocess_image(image_path: str, meta: dict) -> torch.Tensor:
    """Returns a [1, 1, H, W] tensor ready for inference."""
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if image is None:
        raise ValueError(f"Could not read image: {image_path}")

    # Use OpenCV preprocessing instead of albumentations
    image = preprocess_with_cv2(image, image_size=meta["image_size"])

    image = image.astype(np.float32) / 255.0
    image = (image - meta["global_mean"]) / meta["global_std"]
    tensor = torch.from_numpy(image).unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
    return tensor


# ─── Model Loader (singleton) ──────────────────────────────────────────────────
WEIGHTS_PATH = "trainedmodels/Model.pth"
META_PATH    = "trainedmodels/Model.json"
DEVICE       = "cpu"
_model_cache = {}

def load_model(weights_path: str, device: str, meta_path: str = None) -> tuple:
    if weights_path in _model_cache:
        return _model_cache[weights_path]

    # ── Always resolve paths using Path() to normalize slashes ──
    weights_path = Path(weights_path).as_posix()

    if meta_path is None:
        meta_path = Path(weights_path).with_suffix(".json").as_posix()
    else:
        meta_path = Path(meta_path).as_posix()   # normalize whatever is passed in

    if not Path(meta_path).exists():
        raise FileNotFoundError(f"Metadata file not found: {meta_path}")

    with open(meta_path) as f:
        meta = json.load(f)

    model = DenseNet121ForBinaryClassification(
        num_classes=1, pretrained=False, dropout=meta["dropout"]
    )
    state = torch.load(weights_path, map_location=device, weights_only=True)
    model.load_state_dict(state)
    model.to(device)
    model.eval()

    _model_cache[weights_path] = (model, meta)
    return model, meta


# ─── Single Inference ──────────────────────────────────────────────────────────

def run_inference(image_path: str, weights_path: str, device: str = "cpu",
                  generate_gradcam: bool = False) -> dict:
    model, meta = load_model(weights_path, device)
    tensor = preprocess_image(image_path, meta).to(device)

    with torch.no_grad():
        logit = model(tensor)
        prob  = torch.sigmoid(logit).item()

    threshold = meta["best_threshold"]
    pred_class = int(prob >= threshold)
    
    if pred_class == 1:
        label = "Cancer"
        display_prob = prob
    else:
        label = "Normal"
        display_prob = 1 - prob
    
    result = {
        "probability": round(display_prob, 4),
        "predicted_class": pred_class,
        "label": label,
        "threshold_used": threshold,
        "gradcam_overlay": None
    }

    if generate_gradcam:
        gradcam = GradCAMPlusPlus(model, target_layer_name="cbam")
        try:
            # GradCAM needs gradients — don't use no_grad here
            tensor_gc = preprocess_image(image_path, meta).to(device)
            cam_map   = gradcam.generate_cam(tensor_gc, class_idx=pred_class)

            # Build overlay
            orig = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            orig = cv2.resize(orig, (meta["image_size"], meta["image_size"]))
            orig_rgb  = cv2.cvtColor(orig, cv2.COLOR_GRAY2RGB)
            cam_resized = cv2.resize(cam_map, (orig_rgb.shape[1], orig_rgb.shape[0]))

            # Use matplotlib jet colormap for better quality
            import matplotlib.cm as mpl_cm
            heatmap = (mpl_cm.jet(cam_resized)[:, :, :3] * 255).astype(np.uint8)
            overlay = cv2.addWeighted(orig_rgb, 0.6, heatmap, 0.4, 0)
            result["gradcam_overlay"] = overlay   # numpy array, encode downstream
        finally:
            gradcam.cleanup()

    return result