Implement full test-set evaluation with metrics, visualizations, and overlays

- MetricTracker for F1/IoU/Precision/Recall/OA on raw logits
- results.json with all metrics and metadata
- 5x4 prediction grid (Before|After|GT|Pred) + 20 individual plots
- Top-10 overlay images ranked by predicted change area
- Auto eval batch size at 2x training (no gradients needed)
- Colab vs local path resolution, formatted console metrics table

evaluate.py

@@ -1,17 +1,30 @@
-"""Evaluation script for change detection models.
+"""Evaluate a trained change-detection model on the test set.
 
-Runs a trained model on the test set, computes all metrics, and generates
-visualization outputs.
+Computes all metrics (F1, IoU, Precision, Recall, OA), saves a
+``results.json``, generates a 20-sample prediction grid, and produces
+overlay images for the top-10 predictions with the largest predicted
+change area.
 
 Usage:
-    python evaluate.py --config configs/config.yaml --checkpoint checkpoints/unet_pp_best.pth
+    python evaluate.py --config configs/config.yaml \
+        --checkpoint checkpoints/unet_pp_best.pth
+
+    python evaluate.py --config configs/config.yaml \
+        --checkpoint checkpoints/changeformer_best.pth \
+        --model changeformer --output_dir ./my_outputs
 """
 
 import argparse
+import json
 import logging
+import time
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, List, Tuple
 
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
@@ -20,115 +33,397 @@ import yaml
 
 from data.dataset import ChangeDetectionDataset
 from models import get_model
-from utils.metrics import ConfusionMatrix
-from utils.visualization import plot_prediction
+from utils.metrics import MetricTracker
+from utils.visualization import overlay_changes, plot_prediction
 
 logger = logging.getLogger(__name__)
 
 
-def evaluate(
+# ---------------------------------------------------------------------------
+# GPU / batch-size helpers
+# ---------------------------------------------------------------------------
+
+def _detect_gpu_type() -> str:
+    """Detect the current GPU type for batch-size selection.
+
+    Returns:
+        One of ``'T4'``, ``'V100'``, or ``'default'``.
+    """
+    if not torch.cuda.is_available():
+        return "default"
+    name = torch.cuda.get_device_name(0).upper()
+    if "T4" in name:
+        return "T4"
+    elif "V100" in name:
+        return "V100"
+    return "default"
+
+
+def get_train_batch_size(config: Dict[str, Any], model_name: str) -> int:
+    """Look up the *training* batch size for the current GPU + model.
+
+    Args:
+        config: Full project config dict.
+        model_name: Model identifier string.
+
+    Returns:
+        Training batch size as an integer.
+    """
+    gpu_type = _detect_gpu_type()
+    model_sizes = config.get("batch_sizes", {}).get(model_name, {})
+    return model_sizes.get(gpu_type, model_sizes.get("default", 4))
+
+
+# ---------------------------------------------------------------------------
+# Path resolution (same logic as train.py)
+# ---------------------------------------------------------------------------
+
+def resolve_paths(config: Dict[str, Any]) -> Dict[str, Path]:
+    """Build a path dict based on whether Colab mode is enabled.
+
+    Args:
+        config: Full project config dict.
+
+    Returns:
+        Dict with keys ``'data'``, ``'checkpoints'``, ``'logs'``,
+        ``'outputs'``.
+    """
+    if config.get("colab", {}).get("enabled", False):
+        c = config["colab"]
+        return {
+            "data": Path(c["data_dir"]),
+            "checkpoints": Path(c["checkpoint_dir"]),
+            "logs": Path(c["log_dir"]),
+            "outputs": Path(c["output_dir"]),
+        }
+
+    p = config.get("paths", {})
+    return {
+        "data": Path(p.get("processed_data", "./processed_data")),
+        "checkpoints": Path(p.get("checkpoint_dir", "./checkpoints")),
+        "logs": Path(p.get("log_dir", "./logs")),
+        "outputs": Path(p.get("output_dir", "./outputs")),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Evaluation pass
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def run_evaluation(
     model: nn.Module,
     loader: DataLoader,
     device: torch.device,
-    threshold: float = 0.5,
-    output_dir: Path = Path("./outputs"),
-    max_vis: int = 20,
-) -> Dict[str, float]:
-    """Evaluate model on the full test set.
+    tracker: MetricTracker,
+) -> Tuple[Dict[str, float], List[Dict[str, torch.Tensor]]]:
+    """Run inference on the full test set and collect per-sample data.
 
     Args:
-        model: Trained change detection model.
-        loader: Test DataLoader.
+        model: Trained change-detection model (set to eval internally).
+        loader: Test ``DataLoader``.
         device: Target device.
-        threshold: Binarization threshold for predictions.
-        output_dir: Directory to save visualization outputs.
-        max_vis: Maximum number of sample predictions to save.
+        tracker: ``MetricTracker`` (reset externally before this call).
 
     Returns:
-        Dict of metric name -> value.
+        Tuple of ``(metrics_dict, samples_list)``.
+        Each entry in ``samples_list`` is a dict with keys
+        ``'A'``, ``'B'``, ``'mask'``, ``'pred'``, ``'change_area'``
+        (all single-sample tensors on CPU).
     """
     model.eval()
-    cm = ConfusionMatrix()
-    vis_dir = output_dir / "visualizations"
-    vis_dir.mkdir(parents=True, exist_ok=True)
-    vis_count = 0
+    all_samples: List[Dict[str, Any]] = []
+
+    for batch in tqdm(loader, desc="Evaluating", dynamic_ncols=True):
+        img_a = batch["A"].to(device, non_blocking=True)
+        img_b = batch["B"].to(device, non_blocking=True)
+        mask = batch["mask"].to(device, non_blocking=True)
+
+        logits = model(img_a, img_b)
+        tracker.update(logits, mask)
+
+        preds = (torch.sigmoid(logits) >= tracker.threshold).float()
+
+        # Store each sample for later visualisation / ranking
+        for i in range(img_a.size(0)):
+            pred_i = preds[i].cpu()
+            change_area = pred_i.sum().item()
+            all_samples.append({
+                "A": img_a[i].cpu(),
+                "B": img_b[i].cpu(),
+                "mask": mask[i].cpu(),
+                "pred": pred_i,
+                "change_area": change_area,
+            })
+
+    metrics = tracker.compute()
+    return metrics, all_samples
+
+
+# ---------------------------------------------------------------------------
+# Visualisation helpers
+# ---------------------------------------------------------------------------
+
+def save_prediction_grid(
+    samples: List[Dict[str, torch.Tensor]],
+    save_path: Path,
+    num_rows: int = 5,
+) -> None:
+    """Save a grid of sample predictions (Before | After | GT | Pred).
+
+    Args:
+        samples: List of per-sample dicts from ``run_evaluation``.
+        save_path: Destination image path.
+        num_rows: Number of rows in the grid (4 columns each).
+    """
+    num_samples = min(num_rows, len(samples))
+    fig, axes = plt.subplots(num_samples, 4, figsize=(16, 4 * num_samples))
+
+    if num_samples == 1:
+        axes = axes[np.newaxis, :]
+
+    from utils.visualization import _denorm_tensor, _mask_to_numpy
+
+    col_titles = ["Before (A)", "After (B)", "Ground Truth", "Prediction"]
+
+    for row in range(num_samples):
+        s = samples[row]
+        images = [
+            _denorm_tensor(s["A"]),
+            _denorm_tensor(s["B"]),
+            _mask_to_numpy(s["mask"]),
+            (_mask_to_numpy(s["pred"]) > 0.5).astype(np.float32),
+        ]
+        cmaps = [None, None, "gray", "gray"]
+
+        for col in range(4):
+            ax = axes[row, col]
+            ax.imshow(images[col], cmap=cmaps[col], vmin=0, vmax=1)
+            ax.axis("off")
+            if row == 0:
+                ax.set_title(col_titles[col], fontsize=12)
 
-    with torch.no_grad():
-        for batch in tqdm(loader, desc="Evaluating"):
-            img_a = batch["A"].to(device)
-            img_b = batch["B"].to(device)
-            mask = batch["mask"].to(device)
+    fig.tight_layout(pad=1.0)
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(save_path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    logger.info("Saved prediction grid (%d samples): %s", num_samples, save_path)
 
-            logits = model(img_a, img_b)
-            preds = (torch.sigmoid(logits) > threshold).float()
-            cm.update(preds, mask)
 
-            # Save sample visualizations
-            if vis_count < max_vis:
-                for i in range(min(img_a.size(0), max_vis - vis_count)):
-                    plot_prediction(
-                        img_a[i], img_b[i], mask[i], preds[i],
-                        save_path=vis_dir / f"sample_{vis_count:04d}.png",
-                    )
-                    vis_count += 1
+def save_top_overlays(
+    samples: List[Dict[str, torch.Tensor]],
+    output_dir: Path,
+    top_k: int = 10,
+) -> None:
+    """Save overlay images for the top-K predictions by predicted change area.
 
-    metrics = cm.compute()
-    return metrics
+    Args:
+        samples: List of per-sample dicts from ``run_evaluation``.
+        output_dir: Directory to save overlay PNGs.
+        top_k: Number of overlays to save.
+    """
+    import cv2
+
+    overlay_dir = output_dir / "overlays"
+    overlay_dir.mkdir(parents=True, exist_ok=True)
+
+    # Sort by predicted change area (descending) — most "interesting" first
+    ranked = sorted(samples, key=lambda s: s["change_area"], reverse=True)
+    num = min(top_k, len(ranked))
+
+    for idx in range(num):
+        s = ranked[idx]
+        overlay_img = overlay_changes(
+            img_after=s["B"],
+            mask_pred=s["pred"],
+            alpha=0.4,
+            color=(255, 0, 0),
+        )
+        save_file = overlay_dir / f"top_{idx + 1:02d}_area_{s['change_area']:.0f}.png"
+        cv2.imwrite(str(save_file), cv2.cvtColor(overlay_img, cv2.COLOR_RGB2BGR))
+
+    logger.info("Saved %d overlay images: %s", num, overlay_dir)
 
 
+# ---------------------------------------------------------------------------
+# Console formatting
+# ---------------------------------------------------------------------------
+
+def print_metrics_table(
+    metrics: Dict[str, float],
+    model_name: str,
+    checkpoint_path: Path,
+    epoch: int,
+) -> None:
+    """Print a formatted metrics table to the console.
+
+    Args:
+        metrics: Dict of metric name to value.
+        model_name: Model architecture name.
+        checkpoint_path: Path to the loaded checkpoint.
+        epoch: Training epoch the checkpoint was saved at.
+    """
+    border = "=" * 50
+    logger.info(border)
+    logger.info("  TEST SET RESULTS")
+    logger.info(border)
+    logger.info("  Model      : %s", model_name)
+    logger.info("  Checkpoint : %s", checkpoint_path)
+    logger.info("  Epoch      : %d", epoch)
+    logger.info(border)
+    logger.info("  %-12s  %s", "METRIC", "VALUE")
+    logger.info("  " + "-" * 24)
+    for name, value in metrics.items():
+        logger.info("  %-12s  %.4f", name.upper(), value)
+    logger.info(border)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
 def main() -> None:
-    """Main evaluation entry point."""
-    parser = argparse.ArgumentParser(description="Evaluate change detection model")
-    parser.add_argument("--config", type=Path, default=Path("configs/config.yaml"))
-    parser.add_argument("--checkpoint", type=Path, required=True)
-    parser.add_argument("--model", type=str, default=None, help="Override model name")
-    parser.add_argument("--threshold", type=float, default=None)
+    """Entry point — parse CLI args, evaluate model, save outputs."""
+    parser = argparse.ArgumentParser(
+        description="Evaluate a trained change-detection model on the test set",
+    )
+    parser.add_argument(
+        "--config", type=Path, default=Path("configs/config.yaml"),
+        help="Path to the YAML configuration file.",
+    )
+    parser.add_argument(
+        "--checkpoint", type=Path, required=True,
+        help="Path to the model checkpoint (.pth).",
+    )
+    parser.add_argument(
+        "--model", type=str, default=None,
+        help="Override the model name from config.",
+    )
+    parser.add_argument(
+        "--output_dir", type=Path, default=None,
+        help="Override the output directory (default: from config).",
+    )
+    parser.add_argument(
+        "--threshold", type=float, default=None,
+        help="Override the binarisation threshold (default: from config).",
+    )
     args = parser.parse_args()
 
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
 
-    with open(args.config, "r") as f:
-        config = yaml.safe_load(f)
+    # ---- Config -------------------------------------------------------
+    with open(args.config, "r") as fh:
+        config: Dict[str, Any] = yaml.safe_load(fh)
+
+    model_name: str = args.model or config["model"]["name"]
+    threshold: float = args.threshold or config.get("evaluation", {}).get("threshold", 0.5)
 
-    model_name = args.model or config["model"]["name"]
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    threshold = args.threshold or config.get("evaluation", {}).get("threshold", 0.5)
-
-    # Resolve paths
-    colab = config.get("colab", {})
-    if colab.get("enabled", False):
-        data_dir = Path(colab["data_dir"])
-        output_dir = Path(colab["output_dir"])
-    else:
-        data_dir = Path(config["paths"]["processed_data"])
-        output_dir = Path(config["paths"]["output_dir"])
-
-    # Model
+    logger.info("Device: %s", device)
+
+    # ---- Paths --------------------------------------------------------
+    paths = resolve_paths(config)
+    output_dir = args.output_dir or paths["outputs"]
+    output_dir = Path(output_dir) / model_name
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # ---- Load model ---------------------------------------------------
     model = get_model(model_name, config).to(device)
     ckpt = torch.load(args.checkpoint, map_location=device)
     model.load_state_dict(ckpt["model_state_dict"])
-    logger.info("Loaded checkpoint: %s (epoch %d, F1 %.4f)",
-                args.checkpoint, ckpt.get("epoch", -1), ckpt.get("best_f1", -1))
 
-    # Test data
+    ckpt_epoch = ckpt.get("epoch", -1)
+    ckpt_f1 = ckpt.get("best_f1", -1.0)
+    logger.info(
+        "Loaded checkpoint: %s (epoch %d, best F1 %.4f)",
+        args.checkpoint, ckpt_epoch, ckpt_f1,
+    )
+
+    param_count = sum(p.numel() for p in model.parameters()) / 1e6
+    logger.info("Model: %s (%.2fM parameters)", model_name, param_count)
+
+    # ---- Test data ----------------------------------------------------
+    # No gradients stored during eval → safe to use 2x training batch size
+    train_bs = get_train_batch_size(config, model_name)
+    eval_bs = train_bs * 2
+
     ds_cfg = config.get("dataset", {})
-    test_ds = ChangeDetectionDataset(data_dir / "test", split="test", config=config)
+    test_ds = ChangeDetectionDataset(
+        root=paths["data"] / "test", split="test", config=config,
+    )
     test_loader = DataLoader(
-        test_ds, batch_size=8, shuffle=False,
+        test_ds,
+        batch_size=eval_bs,
+        shuffle=False,
         num_workers=ds_cfg.get("num_workers", 4),
         pin_memory=ds_cfg.get("pin_memory", True),
     )
+    logger.info(
+        "Test set: %d samples, %d batches (batch_size=%d, 2x train)",
+        len(test_ds), len(test_loader), eval_bs,
+    )
 
-    # Evaluate
-    metrics = evaluate(model, test_loader, device, threshold, output_dir)
+    # ---- Run evaluation -----------------------------------------------
+    tracker = MetricTracker(threshold=threshold)
+    wall_start = time.monotonic()
 
-    # Print results
-    logger.info("=" * 50)
-    logger.info("TEST SET RESULTS — %s", model_name)
-    logger.info("=" * 50)
-    for name, value in metrics.items():
-        logger.info("  %-12s: %.4f", name.upper(), value)
-    logger.info("=" * 50)
+    metrics, all_samples = run_evaluation(model, test_loader, device, tracker)
+
+    eval_time = time.monotonic() - wall_start
+    logger.info("Evaluation completed in %.1fs", eval_time)
+
+    # ---- Print formatted table ----------------------------------------
+    print_metrics_table(metrics, model_name, args.checkpoint, ckpt_epoch)
+
+    # ---- Save results.json --------------------------------------------
+    results = {
+        "model": model_name,
+        "checkpoint": str(args.checkpoint),
+        "epoch": ckpt_epoch,
+        "threshold": threshold,
+        "num_test_samples": len(test_ds),
+        "eval_time_seconds": round(eval_time, 2),
+        "metrics": {k: round(v, 6) for k, v in metrics.items()},
+    }
+    results_path = output_dir / "results.json"
+    with open(results_path, "w") as f:
+        json.dump(results, f, indent=2)
+    logger.info("Saved results: %s", results_path)
+
+    # ---- Prediction grid (20 samples, 5 rows x 4 cols) ----------------
+    save_prediction_grid(
+        samples=all_samples,
+        save_path=output_dir / "prediction_grid.png",
+        num_rows=min(5, len(all_samples)),
+    )
+
+    # ---- Individual sample plots (up to 20) ---------------------------
+    vis_dir = output_dir / "predictions"
+    vis_dir.mkdir(parents=True, exist_ok=True)
+    num_individual = min(20, len(all_samples))
+    for idx in range(num_individual):
+        s = all_samples[idx]
+        plot_prediction(
+            img_a=s["A"],
+            img_b=s["B"],
+            mask_true=s["mask"],
+            mask_pred=s["pred"],
+            filename=vis_dir / f"sample_{idx + 1:03d}.png",
+        )
+    logger.info("Saved %d individual prediction plots: %s", num_individual, vis_dir)
+
+    # ---- Top-10 overlay images (by predicted change area) -------------
+    save_top_overlays(
+        samples=all_samples,
+        output_dir=output_dir,
+        top_k=10,
+    )
+
+    logger.info("All outputs saved to: %s", output_dir)
 
 
 if __name__ == "__main__":