experiment-2107/scripts/confusion_report.py

"""One-off: load the five saved coarse checkpoints, run them on the test split, and print
confusion matrices + top confusions per class. Not part of the training pipeline."""

from pathlib import Path
import numpy as np
import torch
from sklearn.metrics import confusion_matrix

import labels as labels_mod
from model_transformers import SkatingActionClassifier
from model_transformer_bilstm import SkatingTransformerBiLSTMClassifier
from model_gcn import SkatingGCNClassifier

DATA_DIR = Path("dataset_skeletons")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

CHECKPOINTS = {
    "Conv+Transformer 3L": ("model_transformer_coarse.pt", SkatingActionClassifier, dict(num_layers=3)),
    "Conv+Transformer 6L": ("model_transformer_coarse_6L.pt", SkatingActionClassifier, dict(num_layers=6)),
    "Transformer-BiLSTM 3L": ("model_transformer_bilstm_coarse.pt", SkatingTransformerBiLSTMClassifier, dict(num_layers=3)),
    "Transformer-BiLSTM 6L": ("model_transformer_bilstm_coarse_6L.pt", SkatingTransformerBiLSTMClassifier, dict(num_layers=6)),
    "GCN+Transformer 3L": ("model_gcn_coarse.pt", SkatingGCNClassifier, dict(num_layers=3)),
}


def load_test():
    import pickle
    X = np.asarray(pickle.load(open(DATA_DIR / "test_features.pkl", "rb")), dtype=np.float32)
    y = np.asarray(pickle.load(open(DATA_DIR / "test_labels.pkl", "rb")), dtype=np.int64)
    y_coarse = np.array([labels_mod.FINE_TO_COARSE_IDX[int(v)] for v in y], dtype=np.int64)
    return X, y_coarse


@torch.no_grad()
def predict(model, X):
    model.eval()
    out = []
    for i in range(0, len(X), 256):
        xb = torch.from_numpy(X[i : i + 256]).to(DEVICE)
        out.append(model(xb).argmax(1).cpu().numpy())
    return np.concatenate(out)


def main():
    Xte, yte = load_test()
    class_names = [labels_mod.COARSE_TAXONOMY[i] for i in range(len(labels_mod.COARSE_TAXONOMY))]

    all_cms = {}
    for name, (ckpt_name, cls, kwargs) in CHECKPOINTS.items():
        ckpt = torch.load(DATA_DIR / ckpt_name, map_location=DEVICE, weights_only=False)
        mu, sd = ckpt["feature_mean"], ckpt["feature_std"]
        Xte_std = (Xte - mu) / sd
        if cls is SkatingGCNClassifier:
            model = cls(ckpt["num_classes"], **kwargs).to(DEVICE)
        else:
            model = cls(ckpt["in_features"], ckpt["num_classes"], **kwargs).to(DEVICE)
        model.load_state_dict(ckpt["state_dict"])
        preds = predict(model, Xte_std)
        cm = confusion_matrix(yte, preds, labels=list(range(len(class_names))))
        all_cms[name] = cm

    print(f"\nclasses: {class_names}\n")
    for name, cm in all_cms.items():
        print(f"\n{'='*80}\n{name}\n{'='*80}")
        header = "true\\pred".ljust(12) + "".join(n[:6].rjust(7) for n in class_names)
        print(header)
        for i, row in enumerate(cm):
            print(class_names[i].ljust(12) + "".join(str(v).rjust(7) for v in row))
        print("\ntop confusion per true class (excl. correct diagonal):")
        for i, row in enumerate(cm):
            support = row.sum()
            if support == 0:
                continue
            off_diag = row.copy()
            off_diag[i] = -1
            j = off_diag.argmax()
            if off_diag[j] > 0:
                print(f"  {class_names[i]:<12} (n={support:3d}) -> most confused with "
                      f"{class_names[j]:<12} ({off_diag[j]}/{support} = {off_diag[j]/support:.0%})")

    print(f"\n\n{'='*80}\nCROSS-MODEL: which classes does EVERY model struggle with?\n{'='*80}")
    for i, cname in enumerate(class_names):
        accs = []
        for name, cm in all_cms.items():
            support = cm[i].sum()
            acc = cm[i, i] / support if support else float("nan")
            accs.append(acc)
        print(f"  {cname:<12} per-model recall: " + " | ".join(f"{name}={a:.2f}" for name, a in zip(all_cms, accs)))


if __name__ == "__main__":
    main()