evaluate.py

"""Entrypoint de avaliacao no hold-out de teste (+ validacao).

Produz (para TEST e -- por defeito -- tambem VAL):
  - Metricas (accuracy, P/R/F1 por classe, F1 macro, AUC-ROC)
  - Matriz de confusao (figura + CSV)
  - CSV de predicoes por imagem (usado pelo ensemble.py)
  - [Opcional] exemplos de Grad-CAM (apenas test)
  - [Opcional] Test-Time Augmentation hflip (test e val)

Uso:
    python evaluate.py --config configs/baseline.yaml \\
        --checkpoint outputs/checkpoints/baseline_resnet50/best.pt
    python evaluate.py --config ... --checkpoint ... --gradcam
    python evaluate.py --config ... --checkpoint ... --no_tta
    python evaluate.py --config ... --checkpoint ... --skip_val   # so' test
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay

from src.data.loaders import build_dataloaders
from src.models.builder import build_model
from src.training.losses import build_loss
from src.training.metrics import compute_metrics, sklearn_text_report
from src.utils.config import load_config
from src.utils.logging_utils import setup_logger
from src.utils.seed import set_seed


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--config", type=str, required=True)
    p.add_argument("--checkpoint", type=str, required=True,
                   help="Caminho para .pt produzido pelo train.py.")
    p.add_argument("--gradcam", action="store_true",
                   help="Produzir exemplos de Grad-CAM.")
    p.add_argument("--gradcam_n", type=int, default=8)
    p.add_argument("--data_root", type=str, default=None,
                   help="Override do data.root (ex: database/splits/seed1).")
    p.add_argument("--run_name", type=str, default=None,
                   help="Override do cfg.run_name (afeta pastas de output).")
    p.add_argument("--out_dir", type=str, default="outputs",
                   help="Pasta raiz de outputs (default: outputs).")
    # TTA: por defeito OFF (anatomia da ERCP e' orientada; hflip costuma piorar
    # e nao reproduz o setup da Monica/paper). Quem quiser usa --tta.
    p.add_argument("--tta", dest="tta", action="store_true", default=False,
                   help="Test-Time Augmentation com hflip (default: OFF).")
    p.add_argument("--no_tta", dest="tta", action="store_false",
                   help="Desativa TTA (redundante; e' o default).")
    p.add_argument("--skip_val", action="store_true",
                   help="Salta avaliacao no val (default: avalia val e test).")
    return p.parse_args()


def plot_confusion(cm: np.ndarray, class_names: list[str], out_path: Path) -> None:
    """Matriz de confusao: contagens + normalizada por linha (recall)."""
    cm_norm = cm.astype(float) / np.maximum(cm.sum(axis=1, keepdims=True), 1)
    fig, axes = plt.subplots(1, 2, figsize=(13, 5))
    ConfusionMatrixDisplay(cm, display_labels=class_names).plot(
        ax=axes[0], cmap="Blues", colorbar=False, values_format="d",
    )
    axes[0].set_title("Contagens")
    ConfusionMatrixDisplay(cm_norm, display_labels=class_names).plot(
        ax=axes[1], cmap="Blues", colorbar=False, values_format=".2f",
    )
    axes[1].set_title("Normalizada (recall)")
    for ax in axes:
        ax.set_xticklabels(ax.get_xticklabels(), rotation=30, ha="right")
    plt.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def save_predictions_csv(paths: list[str], y_true: np.ndarray, y_pred: np.ndarray,
                          y_probs: np.ndarray, class_names: list[str],
                          out_path: Path) -> None:
    """CSV usado depois pelo ensemble.py."""
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["path", "true_class", "pred_class"]
                        + [f"p_{c}" for c in class_names])
        for i, p in enumerate(paths):
            writer.writerow(
                [p, class_names[int(y_true[i])], class_names[int(y_pred[i])]]
                + [f"{v:.6f}" for v in y_probs[i].tolist()]
            )


@torch.no_grad()
def predict(model, loader, device, loss_fn=None, use_tta: bool = False):
    """Forward em todo o loader. Com TTA: media de logits original + hflip."""
    model.eval()
    all_logits, all_labels, all_paths = [], [], []
    total_loss, n = 0.0, 0
    for batch in loader:
        if len(batch) == 3:
            images, labels, paths = batch
            all_paths.extend(list(paths))
        else:
            images, labels = batch
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        logits = model(images)
        if use_tta:
            logits_flip = model(torch.flip(images, dims=[3]))
            logits = (logits + logits_flip) / 2.0

        if loss_fn is not None:
            loss = loss_fn(logits, labels)
            total_loss += loss.item() * labels.size(0)
            n += labels.size(0)
        all_logits.append(logits.cpu().numpy())
        all_labels.append(labels.cpu().numpy())

    logits = np.concatenate(all_logits, axis=0)
    labels = np.concatenate(all_labels, axis=0)
    preds = logits.argmax(axis=1)
    z = logits - logits.max(axis=1, keepdims=True)
    e = np.exp(z)
    probs = e / e.sum(axis=1, keepdims=True)
    avg_loss = total_loss / max(n, 1) if loss_fn is not None else None
    return preds, labels, probs, all_paths, avg_loss


def run_gradcam_examples(model, cfg, device, out_dir, n=8):
    """Gera figuras Grad-CAM para n imagens (uma por classe quando possivel)."""
    import cv2
    from src.interpretability.gradcam import (
        make_cam_from_cfg, overlay_cam, save_gradcam_figure,
    )

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    loaders = build_dataloaders(cfg, splits=("test",), return_path_test=True)
    loader = loaders["test"]

    cam = make_cam_from_cfg(model, cfg, method="gradcam")
    class_names = list(cfg.data.classes)
    per_class_quota = max(1, n // len(class_names))
    counts = {c: 0 for c in class_names}
    picked = 0

    for batch in loader:
        images, labels, paths = batch
        images = images.to(device)
        with torch.no_grad():
            preds = model(images).argmax(dim=1).cpu().numpy()
        labels_np = labels.numpy()

        for i in range(images.size(0)):
            cls = class_names[labels_np[i]]
            if counts[cls] >= per_class_quota:
                continue
            orig_bgr = cv2.imread(paths[i], cv2.IMREAD_COLOR)
            orig_rgb = cv2.cvtColor(orig_bgr, cv2.COLOR_BGR2RGB)
            h = int(cfg.data.img_size)
            orig_rgb_resized = cv2.resize(orig_rgb, (h, h)).astype(np.float32) / 255.0

            overlay = overlay_cam(
                cam, input_tensor=images[i:i+1],
                original_image_rgb=orig_rgb_resized,
                target_class=int(preds[i]),
            )
            tag = "OK" if preds[i] == labels_np[i] else "ERR"
            fname = (f"{cls}_{tag}_pred-{class_names[preds[i]]}"
                     f"_{Path(paths[i]).stem}.png")
            save_gradcam_figure(overlay, out_dir / fname)
            counts[cls] += 1
            picked += 1
            if picked >= n:
                return


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    if args.data_root is not None:
        cfg.data["root"] = args.data_root
    if args.run_name is not None:
        cfg["run_name"] = args.run_name
    set_seed(int(cfg.get("seed", 42)), deterministic=False)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    run_name = cfg.run_name
    out_dir = ROOT / args.out_dir / "reports" / run_name
    fig_dir = ROOT / args.out_dir / "figures" / run_name
    out_dir.mkdir(parents=True, exist_ok=True)
    fig_dir.mkdir(parents=True, exist_ok=True)

    logger = setup_logger("eval", log_file=out_dir / "eval.log")

    # Modelo
    model = build_model(cfg)
    state = torch.load(args.checkpoint, map_location=device, weights_only=False)
    if "model_state_dict" in state:
        model.load_state_dict(state["model_state_dict"])
    else:
        model.load_state_dict(state)
    model.to(device)
    logger.info(f"Checkpoint: {args.checkpoint}")
    logger.info(f"  best epoch: {state.get('epoch')}, score: {state.get('score')}")
    if state.get("use_ema", False):
        logger.info("  (pesos EMA — usados para inferencia)")
    logger.info(f"  TTA: {'ON (hflip)' if args.tta else 'OFF'}")

    # Loaders: test (com paths para predictions.csv) e val (se nao --skip_val)
    splits = ("val", "test") if not args.skip_val else ("test",)
    loaders = build_dataloaders(cfg, splits=splits, return_path_test=True)

    # IMPORTANTE: weighted_ce precisa do tensor de pesos no mesmo device dos logits.
    # nn.CrossEntropyLoss(weight=...) nao faz .to(device) automaticamente.
    loss_fn = build_loss(cfg, class_counts=np.ones(len(cfg.data.classes))).to(device)
    class_names = list(cfg.data.classes)

    def _evaluate_split(split: str) -> dict:
        """Corre predict + metrics num split. Escreve confusion + predictions
        com sufixo do split. Devolve o resumo (dict)."""
        logger.info(f"\n===== {split.upper()} =====")
        preds, labels, probs, paths, avg_loss = predict(
            model, loaders[split], device, loss_fn, use_tta=args.tta,
        )
        report = compute_metrics(
            y_true=labels, y_pred=preds, y_probs=probs,
            class_names=class_names, loss=avg_loss,
        )
        logger.info("\n" + report.summary())
        logger.info("\n" + sklearn_text_report(labels, preds, class_names))

        cm_fig = fig_dir / f"confusion_matrix_{split}.png"
        cm_csv = out_dir / f"confusion_matrix_{split}.csv"
        preds_csv = out_dir / f"predictions_{split}.csv"
        plot_confusion(report.confusion, class_names, cm_fig)
        np.savetxt(cm_csv, report.confusion, fmt="%d", delimiter=",")
        save_predictions_csv(paths, labels, preds, probs, class_names, preds_csv)

        return {
            "accuracy": report.accuracy,
            "f1_macro": report.f1_macro,
            "f1_weighted": report.f1_weighted,
            "auc_ovr": report.auc_ovr,
            "loss": report.loss,
            "per_class": report.per_class,
        }

    per_split: dict[str, dict] = {sp: _evaluate_split(sp) for sp in splits}

    # JSON resumo: campos de TEST no topo (back-compat) + bloco por split
    test_summary = per_split["test"]
    summary = {
        # Campos legacy (test) -- preservados para nao quebrar consumidores
        "accuracy": test_summary["accuracy"],
        "f1_macro": test_summary["f1_macro"],
        "f1_weighted": test_summary["f1_weighted"],
        "auc_ovr": test_summary["auc_ovr"],
        "loss": test_summary["loss"],
        "per_class": test_summary["per_class"],
        "tta": bool(args.tta),
        "baseline_paper_f1_macro": 0.738,
        # Novo: bloco com todos os splits avaliados (test + opcionalmente val)
        "splits": per_split,
    }
    with (out_dir / "metrics.json").open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    # Linha de resumo compacta (util para varrer logs multi-seed)
    val_line = ""
    if "val" in per_split:
        v = per_split["val"]
        val_line = (f"  VAL  | acc={v['accuracy']:.4f}  f1m={v['f1_macro']:.4f}  "
                    f"auc={v['auc_ovr']:.4f}\n")
    t = per_split["test"]
    test_line = (f"  TEST | acc={t['accuracy']:.4f}  f1m={t['f1_macro']:.4f}  "
                 f"auc={t['auc_ovr']:.4f}")
    logger.info(f"\n--- RESUMO ({run_name}) ---\n{val_line}{test_line}")
    logger.info(f"Resultados: {out_dir}")
    logger.info(f"Figuras:    {fig_dir}")

    if args.gradcam:
        logger.info(f"A gerar {args.gradcam_n} exemplos Grad-CAM...")
        run_gradcam_examples(model, cfg, device,
                              out_dir=fig_dir / "gradcam", n=args.gradcam_n)
        logger.info(f"Grad-CAMs em: {fig_dir / 'gradcam'}")


if __name__ == "__main__":
    main()