evaluate_sam3.py

"""Entrypoint de avaliacao no hold-out de teste para o pipeline SAM3 4-canal.

Uso:
    conda activate vila
    python evaluate_sam3.py --config configs/resnet50_sam3_4ch.yaml \\
        --checkpoint outputs/checkpoints/resnet50_sam3_4ch_seed42/best.pt
"""
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_sam3 import build_sam3_dataloaders as build_dataloaders
from src.models.builder_sam3 import build_model_4ch as 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_sam3.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.")
    p.add_argument("--masks_root", type=str, default=None,
                   help="Override do data.masks_root (mascaras SAM3 dos dados novos).")
    p.add_argument("--run_name", type=str, default=None,
                   help="Override do cfg.run_name.")
    p.add_argument("--tta", dest="tta", action="store_true", default=True,
                   help="Test-Time Augmentation com hflip (default: ON).")
    p.add_argument("--no_tta", dest="tta", action="store_false",
                   help="Desativa TTA.")
    return p.parse_args()


def plot_confusion(cm: np.ndarray, class_names: list[str], out_path: Path) -> None:
    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:
    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):
    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):
    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.masks_root is not None:
        cfg.data["masks_root"] = args.masks_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 / "outputs" / "reports" / run_name
    fig_dir = ROOT / "outputs" / "figures" / run_name
    out_dir.mkdir(parents=True, exist_ok=True)
    fig_dir.mkdir(parents=True, exist_ok=True)

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

    # Modelo
    model = build_model(cfg)
    state = torch.load(args.checkpoint, map_location=device, weights_only=False)
    model.load_state_dict(state["model_state_dict"])
    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 = build_dataloaders(cfg, splits=("test",), return_path_test=True)

    loss_fn = build_loss(cfg, class_counts=np.ones(len(cfg.data.classes))).to(device)

    preds, labels, probs, paths, avg_loss = predict(
        model, loaders["test"], device, loss_fn, use_tta=args.tta,
    )

    class_names = list(cfg.data.classes)
    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))

    plot_confusion(report.confusion, class_names, fig_dir / "confusion_matrix.png")
    np.savetxt(out_dir / "confusion_matrix.csv", report.confusion,
                fmt="%d", delimiter=",")
    save_predictions_csv(paths, labels, preds, probs, class_names,
                         out_dir / "predictions.csv")

    summary = {
        "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,
        "tta": bool(args.tta),
        "baseline_paper_f1_macro": 0.738,
    }
    with (out_dir / "metrics.json").open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    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()