src/egg_damage/explainability.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import cv2
import numpy as np
import pandas as pd
import torch
from PIL import Image, ImageOps
from torch import nn

from .augmentations import IMAGENET_MEAN, IMAGENET_STD, build_eval_transform
from .compare_models import load_best_model_record
from .dl_models import create_model, load_torch_checkpoint
from .paths import ensure_dir
from .preprocessing import load_pil_image
from .utils import get_logger


LOGGER = get_logger(__name__)


def find_last_conv_layer(model: nn.Module) -> nn.Module | None:
    last_conv: nn.Module | None = None
    for module in model.modules():
        if isinstance(module, nn.Conv2d):
            last_conv = module
    return last_conv


def denormalize_tensor(image_tensor: torch.Tensor) -> np.ndarray:
    mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
    std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
    image = image_tensor.detach().cpu() * std + mean
    image = image.clamp(0, 1).permute(1, 2, 0).numpy()
    return (image * 255).astype(np.uint8)


def gradcam_overlay(
    model: nn.Module,
    image: str | Path | Image.Image,
    config: dict[str, Any],
    output_path: str | Path | None = None,
    target_class: int | None = None,
    device: torch.device | None = None,
) -> Image.Image:
    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device).eval()
    target_layer = find_last_conv_layer(model)
    if target_layer is None:
        raise ValueError("Grad-CAM is only available for CNN models with Conv2d layers.")

    activation: dict[str, torch.Tensor] = {}

    def forward_hook(_: nn.Module, __: tuple[torch.Tensor, ...], output: torch.Tensor) -> torch.Tensor:
        cloned = output.clone()
        cloned.retain_grad()
        activation["value"] = cloned
        return cloned

    handle_fwd = target_layer.register_forward_hook(forward_hook)
    try:
        pil = load_pil_image(image, mode="RGB")
        transform = build_eval_transform(config)
        tensor = transform(pil).unsqueeze(0).to(device)
        tensor.requires_grad_(True)
        logits = model(tensor)
        class_idx = int(target_class if target_class is not None else torch.argmax(logits, dim=1).item())
        model.zero_grad(set_to_none=True)
        logits[:, class_idx].sum().backward()
        acts_with_grad = activation.get("value")
        if acts_with_grad is None or acts_with_grad.grad is None:
            raise RuntimeError("Grad-CAM hook did not capture activations/gradients.")
        acts = acts_with_grad.detach()[0]
        grads = acts_with_grad.grad.detach()[0]
        weights = grads.mean(dim=(1, 2), keepdim=True)
        cam = (weights * acts).sum(dim=0)
        cam = torch.relu(cam)
        cam -= cam.min()
        cam /= cam.max().clamp(min=1e-8)
        cam_np = cam.detach().cpu().numpy()
        base = denormalize_tensor(tensor[0])
        heatmap = cv2.resize(cam_np, (base.shape[1], base.shape[0]))
        heatmap_uint8 = np.uint8(255 * heatmap)
        color = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
        color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
        overlay = np.uint8(0.55 * base + 0.45 * color)
        out = Image.fromarray(overlay)
        if output_path:
            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            out.save(output_path)
        return out
    finally:
        handle_fwd.remove()


def save_gradcam_examples_for_best(
    config: dict[str, Any],
    splits_df: pd.DataFrame,
    leaderboard_df: pd.DataFrame | None = None,
) -> list[Path]:
    record = (
        leaderboard_df.iloc[0].to_dict()
        if leaderboard_df is not None and not leaderboard_df.empty
        else load_best_model_record(config)
    )
    if record.get("model_type") != "deep_learning":
        LOGGER.info("Best model is not a deep CNN; Grad-CAM generation skipped.")
        return []
    model_path = Path(record["model_path"])
    checkpoint = load_torch_checkpoint(model_path, map_location="cpu")
    family = checkpoint.get("family", "cnn")
    if family != "cnn":
        LOGGER.info("Best deep model is %s; Grad-CAM is skipped for non-CNN models.", family)
        return []
    model = create_model(checkpoint["model_key"], checkpoint.get("config", config), pretrained=False)
    model.load_state_dict(checkpoint["state_dict"])
    test_df = splits_df[splits_df["split"] == "test"].copy().reset_index(drop=True)
    if test_df.empty:
        return []
    count = int(config.get("explainability", {}).get("max_images", 8))
    sample = test_df.groupby("label", group_keys=False).head(max(1, count // 2)).head(count)
    out_dir = ensure_dir(Path(config["paths"]["output_dir"]) / "plots" / "gradcam")
    saved: list[Path] = []
    for idx, row in enumerate(sample.itertuples(index=False), start=1):
        output_path = out_dir / f"{record['model_name']}_gradcam_{idx:02d}.png"
        try:
            gradcam_overlay(model, row.filepath, checkpoint.get("config", config), output_path, target_class=None)
            saved.append(output_path)
        except Exception as exc:
            LOGGER.warning("Failed Grad-CAM for %s: %s", row.filepath, exc)
    return saved


def gradcam_for_checkpoint(
    model_path: str | Path,
    image: str | Path | Image.Image,
    config: dict[str, Any],
    output_path: str | Path | None = None,
) -> Image.Image:
    checkpoint = load_torch_checkpoint(model_path, map_location="cpu")
    if checkpoint.get("family", "cnn") != "cnn":
        raise ValueError("Grad-CAM is only enabled for CNN deep-learning checkpoints.")
    model = create_model(checkpoint["model_key"], checkpoint.get("config", config), pretrained=False)
    model.load_state_dict(checkpoint["state_dict"])
    return gradcam_overlay(model, image, checkpoint.get("config", config), output_path=output_path)