evaluate.py

"""
Evaluate the cross-attention tight-fusion network.

Loads a trained ``CrossAttentionFusionNet`` checkpoint, runs it over
test sequences, dead-reckons relative poses to a world trajectory, and
reports ATE / RTE / mean rotation error using the shared metrics
module. Writes trajectory + per-pair error plots and a vision-IMU
cosine similarity plot per sequence (the closest analogue to the
gated branch's gate plot).
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.transforms.functional as TF
from PIL import Image

_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from fusion_common.dataset import PairedDataset, _PairedSequence  # noqa: E402
from fusion_common.metrics import trajectory_metrics  # noqa: E402
from model import CrossAttentionFusionNet  # noqa: E402

_IMAGENET_MEAN = (0.485, 0.456, 0.406)
_IMAGENET_STD = (0.229, 0.224, 0.225)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Evaluate cross-attention fusion.")
    p.add_argument("--data_root", type=str, required=True)
    p.add_argument("--checkpoint", type=str, required=True)
    p.add_argument("--results_dir", type=str, default="results/cross_attention")
    p.add_argument("--imu_rate", type=float, default=1000.0)
    p.add_argument("--imu_context", type=int, default=100)
    p.add_argument("--img_height", type=int, default=224)
    p.add_argument("--img_width", type=int, default=224)
    p.add_argument("--rte_interval", type=float, default=5.0)
    p.add_argument("--chunk", type=int, default=10)
    p.add_argument("--sequences", type=str, nargs="*", default=None)
    return p.parse_args()


def _load_image(path: str, h: int, w: int) -> torch.Tensor:
    img = Image.open(path).convert("RGB").resize((w, h), Image.BILINEAR)
    return TF.normalize(TF.to_tensor(img), _IMAGENET_MEAN, _IMAGENET_STD)


def _predict_sequence(model, seq, device, chunk, imu_context, img_h, img_w):
    n_pairs = seq.n_cam - 1
    rel_t = np.zeros((n_pairs, 3), dtype=np.float64)
    rel_R = np.zeros((n_pairs, 3, 3), dtype=np.float64)
    cos_log = np.zeros((n_pairs,), dtype=np.float64)
    pos = int(np.searchsorted(seq.frame_imu_indices, imu_context - 1))
    rel_t_offset = pos
    hidden = None
    while pos < n_pairs:
        end = min(pos + chunk, n_pairs)
        T = end - pos
        f0 = torch.empty((1, T, 3, img_h, img_w), dtype=torch.float32)
        f1 = torch.empty_like(f0)
        acc = torch.empty((1, T, imu_context, 3), dtype=torch.float32)
        gyro = torch.empty_like(acc)
        att = torch.empty_like(acc)
        for i in range(T):
            t0 = pos + i; t1 = pos + i + 1
            f0[0, i] = _load_image(seq.frame_files[t0], img_h, img_w)
            f1[0, i] = _load_image(seq.frame_files[t1], img_h, img_w)
            imu_hi = int(seq.frame_imu_indices[t1])
            lo = imu_hi - imu_context + 1
            acc[0, i] = torch.from_numpy(seq.acc[lo : imu_hi + 1])
            gyro[0, i] = torch.from_numpy(seq.gyro[lo : imu_hi + 1])
            att[0, i] = torch.from_numpy(seq.attitude[lo : imu_hi + 1])
        out = model(
            f0.to(device), f1.to(device),
            acc.to(device), gyro.to(device), att.to(device).float(),
            hidden=hidden,
        )
        hidden = (out["hidden"][0].detach(), out["hidden"][1].detach())
        rel_t[pos:end] = out["trans"][0].detach().cpu().numpy()
        rel_R[pos:end] = out["R"][0].detach().cpu().numpy()
        cos_log[pos:end] = out["vis_imu_cos"][0].detach().cpu().numpy()
        pos = end
    rel_R[:rel_t_offset] = np.eye(3)
    return rel_t, rel_R, cos_log


def _integrate(rel_R, rel_t, T0):
    T_world = [T0.copy()]
    for i in range(rel_R.shape[0]):
        T_rel = np.eye(4); T_rel[:3, :3] = rel_R[i]; T_rel[:3, 3] = rel_t[i]
        T_world.append(T_world[-1] @ T_rel)
    return np.stack(T_world, axis=0)


def _plot(name, world, gt, metrics, cos_log, out_dir):
    out_dir.mkdir(parents=True, exist_ok=True)
    pos_pred = world[:, :3, 3]; pos_gt = gt["pos_gt"]

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.plot(pos_gt[:, 0], pos_gt[:, 1], "b-", label="ground truth")
    ax.plot(pos_pred[:, 0], pos_pred[:, 1], "r-", label="cross-attn fusion")
    ax.scatter([pos_gt[0, 0]], [pos_gt[0, 1]], c="green", s=60, label="start", zorder=5)
    ax.set_xlabel("X [m]"); ax.set_ylabel("Y [m]")
    ax.set_title(f"{name} — trajectory")
    ax.axis("equal"); ax.grid(True, alpha=0.3); ax.legend()
    fig.tight_layout()
    fig.savefig(out_dir / f"{name}_trajectory.png", dpi=150); plt.close(fig)

    fig = plt.figure(figsize=(8, 6))
    ax3 = fig.add_subplot(111, projection="3d")
    ax3.plot(pos_gt[:, 0], pos_gt[:, 1], pos_gt[:, 2], "b-", label="ground truth")
    ax3.plot(pos_pred[:, 0], pos_pred[:, 1], pos_pred[:, 2], "r-", label="cross-attn fusion")
    ax3.scatter([pos_gt[0, 0]], [pos_gt[0, 1]], [pos_gt[0, 2]], c="green", s=60, label="start")
    ax3.set_xlabel("X [m]"); ax3.set_ylabel("Y [m]"); ax3.set_zlabel("Z [m]")
    ax3.set_title(f"{name} — trajectory (3D)")
    ax3.legend()
    fig.tight_layout()
    fig.savefig(out_dir / f"{name}_trajectory_3d.png", dpi=150); plt.close(fig)

    for kind, series, ylabel in [
        ("trans", metrics["per_frame_trans"], "Position error [m]"),
        ("rot", metrics["per_frame_rot_deg"], "Rotation error [deg]"),
    ]:
        fig, ax = plt.subplots(figsize=(8, 3))
        ax.plot(series, "r-")
        ax.set_xlabel("Sample index"); ax.set_ylabel(ylabel)
        ax.set_title(f"{name} — per-frame {kind} error")
        ax.grid(True, alpha=0.3)
        fig.tight_layout()
        fig.savefig(out_dir / f"{name}_{kind}_error.png", dpi=150); plt.close(fig)

    fig, ax = plt.subplots(figsize=(8, 3))
    ax.plot(cos_log, "g-")
    ax.set_xlabel("Pair index"); ax.set_ylabel("Cosine similarity (vision vs IMU token)")
    ax.set_ylim([-1.0, 1.0])
    ax.set_title(f"{name} — vision-IMU agreement")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / f"{name}_vis_imu_cos.png", dpi=150); plt.close(fig)


def main() -> None:
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = CrossAttentionFusionNet().to(device).eval()
    state = torch.load(args.checkpoint, map_location=device)
    model.load_state_dict(state.get("model", state))

    seq_names = args.sequences or PairedDataset.list_sequences(args.data_root)
    out_dir = Path(args.results_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"{'sequence':<20} {'ATE [m]':>10} {'RTE [m]':>10} {'rot [deg]':>12} "
          f"{'pairs':>8}")
    print("-" * 64)
    aggregate = {"ate": [], "rte": [], "rot_deg": []}
    with torch.no_grad():
        for name in seq_names:
            seq = _PairedSequence(Path(args.data_root) / name, imu_rate=args.imu_rate)
            rel_t, rel_R, cos_log = _predict_sequence(
                model, seq, device, args.chunk, args.imu_context,
                args.img_height, args.img_width,
            )
            T0 = np.eye(4); T0[:3, :3] = seq.R_cam[0]; T0[:3, 3] = seq.p_cam[0]
            world = _integrate(rel_R, rel_t, T0)
            n = world.shape[0]
            rte_step = max(1, int(round(args.rte_interval / seq.dt_cam)))
            metrics = trajectory_metrics(
                pos_pred=world[:, :3, 3],
                pos_gt=seq.p_cam[:n],
                R_pred=world[:, :3, :3],
                R_gt=seq.R_cam[:n],
                rte_step=rte_step,
                aligned=True,
            )
            print(f"{name:<20} {metrics['ate']:>10.4f} {metrics['rte']:>10.4f} "
                  f"{metrics['rot_deg']:>12.4f} {n:>8d}")
            _plot(name, world, {"pos_gt": seq.p_cam[:n]}, metrics, cos_log, out_dir)
            aggregate["ate"].append(metrics["ate"])
            aggregate["rte"].append(metrics["rte"])
            aggregate["rot_deg"].append(metrics["rot_deg"])

    print("-" * 64)
    print(f"{'mean':<20} {np.mean(aggregate['ate']):>10.4f} "
          f"{np.mean(aggregate['rte']):>10.4f} "
          f"{np.mean(aggregate['rot_deg']):>12.4f}")
    print(f"\nPlots written to {out_dir.resolve()}")


if __name__ == "__main__":
    main()