geospatiallabamsob/mantis-vision-api

1"""Evaluate the best multi-head checkpoint on the held-out test split.2 3Reports, per measurement:4  - classification: accuracy, precision/recall/F1 (macro + per-class),5    confusion matrix image (the schema's primary classification only), and6    one-vs-rest ROC AUC7  - regression: mean absolute error, on the samples where it applies8  - segmentation: mean IoU per mask class, on the samples with a ground-truth9    mask10 11Per spec: never judge the model on accuracy alone.12 13Usage:14    python -m src.evaluate15"""16from __future__ import annotations17 18import json19import sys20from pathlib import Path21 22import matplotlib.pyplot as plt23import numpy as np24import torch25import torch.nn.functional as F26from sklearn.metrics import (27    ConfusionMatrixDisplay,28    classification_report,29    confusion_matrix,30    roc_auc_score,31)32 33sys.path.insert(0, str(Path(__file__).resolve().parents[1]))34from config import Config, Schema, config as _default_config  # noqa: E40235from src.data.dataset import get_dataloaders  # noqa: E40236from src.models.efficientnet import load_checkpoint  # noqa: E40237from src.utils.logger import get_logger  # noqa: E40238from src.utils.seed import get_device, set_seed  # noqa: E40239 40 41def _mae(pred: list[float], target: list[float], mask: list[float]) -> float | None:42    pairs = [(p, t) for p, t, m in zip(pred, target, mask) if m > 0.5]43    if not pairs:44        return None45    return float(np.mean([abs(p - t) for p, t in pairs]))46 47 48def _mean_iou(pred_masks: list[np.ndarray], target_masks: list[np.ndarray], num_classes: int) -> dict[str, float] | None:49    if not pred_masks:50        return None51    ious_per_class: dict[int, list[float]] = {c: [] for c in range(num_classes)}52    for pred, target in zip(pred_masks, target_masks):53        for c in range(num_classes):54            pred_c = pred == c55            target_c = target == c56            union = np.logical_or(pred_c, target_c).sum()57            if union == 0:58                continue59            intersection = np.logical_and(pred_c, target_c).sum()60            ious_per_class[c].append(float(intersection) / float(union))61    return {str(c): (float(np.mean(v)) if v else None) for c, v in ious_per_class.items()}62 63 64def evaluate(checkpoint_path: Path | None = None, cfg: Config | None = None) -> dict:65    cfg = cfg if cfg is not None else _default_config66    set_seed(cfg.seed)67    device = get_device(cfg.device)68    logger = get_logger("evaluate", cfg.logs_dir)69 70    checkpoint_path = checkpoint_path or (cfg.checkpoints_dir / "best_model.pt")71    model, schema = load_checkpoint(checkpoint_path, device)72    logger.info("Loaded checkpoint %s (measurements=%s)", checkpoint_path, [m.key for m in schema.measurements])73 74    data = get_dataloaders(cfg, schema)75 76    results: dict = {}77 78    primary = schema.primary_classification()79 80    with torch.no_grad():81        per_measurement_state: dict[str, dict] = {82            m.key: (83                {"labels": [], "preds": [], "probs": []}84                if m.type == "classification"85                else {"pred": [], "target": [], "mask": []}86                if m.type == "regression"87                else {"pred_masks": [], "target_masks": []}88            )89            for m in schema.measurements90        }91 92        for images, targets in data.test:93            images = images.to(device)94            outputs = model(images)95 96            for m in schema.measurements:97                state = per_measurement_state[m.key]98                if m.type == "classification":99                    mask = targets[f"{m.key}_mask"].numpy()100                    probs = F.softmax(outputs[m.key], dim=1).cpu().numpy()101                    preds = probs.argmax(axis=1)102                    ids = targets[f"{m.key}_id"].numpy()103                    for keep, p, t, prob in zip(mask > 0.5, preds, ids, probs):104                        if keep:105                            state["preds"].append(int(p))106                            state["labels"].append(int(t))107                            state["probs"].append(prob.tolist())108                elif m.type == "regression":109                    state["pred"].extend(outputs[m.key].cpu().numpy().tolist())110                    state["target"].extend(targets[m.key].numpy().tolist())111                    state["mask"].extend(targets[f"{m.key}_mask"].numpy().tolist())112                elif m.type == "segmentation":113                    seg_mask = targets[f"{m.key}_seg_mask"].numpy()114                    pred_classes = outputs[m.key].argmax(dim=1).cpu().numpy()115                    target_classes = targets[f"{m.key}_seg"].numpy()116                    for keep, p, t in zip(seg_mask > 0.5, pred_classes, target_classes):117                        if keep:118                            state["pred_masks"].append(p)119                            state["target_masks"].append(t)120 121    for m in schema.measurements:122        state = per_measurement_state[m.key]123 124        if m.type == "classification":125            class_names = m.class_names()126            if not state["labels"]:127                results[m.key] = None128                continue129            labels_arr = np.array(state["labels"])130            preds_arr = np.array(state["preds"])131            probs_arr = np.array(state["probs"])132 133            report = classification_report(134                labels_arr, preds_arr, labels=list(range(len(class_names))),135                target_names=class_names, output_dict=True, zero_division=0,136            )137            logger.info(138                "[%s]\n%s",139                m.key,140                classification_report(labels_arr, preds_arr, labels=list(range(len(class_names))), target_names=class_names, zero_division=0),141            )142 143            cm = confusion_matrix(labels_arr, preds_arr, labels=list(range(len(class_names))))144            per_class_accuracy = {145                name: float(cm[i, i] / cm[i].sum()) if cm[i].sum() else 0.0 for i, name in enumerate(class_names)146            }147 148            # Per-class one-vs-rest ROC AUC, computed one class at a time by149            # binarizing y_true against that class rather than via a single150            # roc_auc_score(multi_class="ovr", average=None) call.151            #152            # That single-call form is unusable here because its return SHAPE153            # depends on the data, not just the schema: with a small/sparse154            # labeled set (the norm right after an admin adds a measurement and155            # labels a handful of images) a split routinely contains only one156            # of a measurement's classes, and sklearn then returns a bare NaN157            # *scalar* for the 2-declared-classes case (but a NaN *array* for158            # 3+). Iterating that scalar is the 'float' object is not iterable159            # crash. Binarizing per class sidesteps sklearn's shape-guessing160            # entirely: every call gets a clean 1-D y_true and a 1-D score161            # column, so the result is always a plain float we control.162            #163            # A class's OvR AUC is undefined unless BOTH that class and the164            # "rest" appear in y_true; those cases (and any residual sklearn165            # ValueError) map to None. None (not NaN) matters downstream:166            # json.dumps emits NaN as the bare token `NaN`, which PostgREST167            # rejects as invalid JSON (an opaque HTTP 400 when reporting168            # results back to model_runs).169            roc_auc: dict[str, float | None] = {}170            for i, name in enumerate(class_names):171                y_true_binary = (labels_arr == i).astype(int)172                # Needs both a positive and a negative present to be defined.173                if y_true_binary.min() == y_true_binary.max():174                    roc_auc[name] = None175                    continue176                try:177                    score = roc_auc_score(y_true_binary, probs_arr[:, i])178                    roc_auc[name] = float(score) if not np.isnan(score) else None179                except (ValueError, IndexError) as e:180                    logger.warning("[%s] Could not compute ROC AUC for class %r: %s", m.key, name, e)181                    roc_auc[name] = None182 183            cm_path = None184            if m is primary:185                # Only the primary classification gets a confusion-matrix186                # image — the others (e.g. disease_subtype) are reported as187                # plain classification_report dicts.188                cfg.reports_dir.mkdir(parents=True, exist_ok=True)189                fig, ax = plt.subplots(figsize=(8, 7))190                ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names).plot(191                    ax=ax, cmap="Blues", xticks_rotation=45192                )193                plt.tight_layout()194                cm_path = cfg.reports_dir / "confusion_matrix.png"195                fig.savefig(cm_path, dpi=150)196                plt.close(fig)197 198            results[m.key] = {199                "accuracy": report["accuracy"],200                "macro_avg": report["macro avg"],201                "weighted_avg": report["weighted avg"],202                "per_class": {203                    name: {204                        "precision": report[name]["precision"],205                        "recall": report[name]["recall"],206                        "f1_score": report[name]["f1-score"],207                        "support": report[name]["support"],208                        "accuracy": per_class_accuracy[name],209                        "roc_auc": roc_auc.get(name),210                    }211                    for name in class_names212                },213                "confusion_matrix_path": str(cm_path) if cm_path else None,214            }215 216        elif m.type == "regression":217            results[m.key] = {"mae": _mae(state["pred"], state["target"], state["mask"])}218 219        elif m.type == "segmentation":220            results[m.key] = {"mean_iou_per_class": _mean_iou(state["pred_masks"], state["target_masks"], len(m.seg_classes))}221 222    results_path = cfg.reports_dir / "evaluation_results.json"223    cfg.reports_dir.mkdir(parents=True, exist_ok=True)224    results_path.write_text(json.dumps(results, indent=2))225    logger.info("Saved evaluation results -> %s", results_path)226 227    return results228 229 230if __name__ == "__main__":231    evaluate()232