test.py

import os
import argparse
import tqdm
from pathlib import Path

import torch
from torch.utils.data import DataLoader, TensorDataset

from utils import (
    DATASET_CHOICES,
    set_seed,
    load_backbone,
    load_dataset_config,
    load_test_embeddings,
    get_text_embs,
    spearman_corr_from_logits,
    kendall_tau_topk_from_logits,
    kl_div_from_logits,
    backbone_to_name
)

def get_dataloader(seen_embs, seen_targets, unseen_embs, unseen_targets, batch_size):
    seen_loader = DataLoader(TensorDataset(seen_embs, seen_targets), batch_size=batch_size, num_workers=0)
    unseen_loader = DataLoader(TensorDataset(unseen_embs, unseen_targets), batch_size=batch_size, num_workers=0)
    return seen_loader, unseen_loader

def test(val_loader, class_name_embs, target_mapping, A_matrix, device, kendall_topk=50):
    clip_true_count = 0
    ezpc_true_count = 0
    total_count = 0

    agree_sum = 0.0
    spearman_sum = 0.0
    kendall_sum = 0.0
    kl_sum = 0.0

    for batch_embs, batch_ids in tqdm.tqdm(val_loader, desc="Testing"):
        batch_embs = batch_embs.to(device)
        if target_mapping is not None:
            batch_ids = torch.tensor([target_mapping[idx.item()] for idx in batch_ids], device=device)
        else:
            batch_ids = batch_ids.to(device)

        clip_logits = batch_embs @ class_name_embs.T
        ezpc_logits = batch_embs @ A_matrix @ A_matrix.T @ class_name_embs.T

        clip_preds = clip_logits.argmax(dim=1)
        ezpc_preds = ezpc_logits.argmax(dim=1)

        clip_true_count += (clip_preds == batch_ids).sum().item()
        ezpc_true_count += (ezpc_preds == batch_ids).sum().item()
        total_count += batch_embs.size(0)

        # Fidelity metrics
        agree_sum += (clip_preds == ezpc_preds).float().sum().item()
        spearman_sum += spearman_corr_from_logits(clip_logits, ezpc_logits).sum().item()

        if kendall_topk is not None and kendall_topk > 1:
            kendall_sum += kendall_tau_topk_from_logits(clip_logits, ezpc_logits, topk=kendall_topk).sum().item()

        kl_sum += kl_div_from_logits(clip_logits, ezpc_logits).sum().item()

    clip_acc = clip_true_count / total_count
    ezpc_acc = ezpc_true_count / total_count

    top1_agree = 100.0 * agree_sum / total_count
    spearman = spearman_sum / total_count
    kendall = (kendall_sum / total_count) if (kendall_topk is not None and kendall_topk > 1) else float("nan")
    kl = kl_sum / total_count

    return clip_acc, ezpc_acc, top1_agree, spearman, kendall, kl

def _resolve_text_embs_path(dataset_root, dataset, backbone, kind):
    """Return the path to cached text embeddings in the dataset embeddings dir, or None.

    Text embeddings depend only on (dataset, backbone), so they live alongside the
    image embeddings and are generated by data/save_text_embs.py.
      kind='classname' -> {backbone}_classname_embs.pt
      kind='concept'   -> {backbone}_concept_matrix.pt
    """
    bname = backbone_to_name(backbone)
    fname = f"{bname}_classname_embs.pt" if kind == "classname" else f"{bname}_concept_matrix.pt"
    path = Path(dataset_root) / dataset / "embeddings" / fname
    return path if path.exists() else None


def _load_or_compute_text_embs(model, str_list, args, tokenizer, kind="classname"):
    """Load cached text embeddings from the dataset embeddings dir, else compute fresh."""
    cache_path = (
        None if args.recompute_text_embs
        else _resolve_text_embs_path(args.dataset_root, args.dataset, args.backbone, kind)
    )
    if cache_path is not None:
        print(f"Loading cached text embeddings from {cache_path}")
        return torch.load(cache_path, map_location=args.device, weights_only=True)
    print("Computing text embeddings on the fly")
    return get_text_embs(model, str_list, args.backbone, tokenizer, args.device)


def main(args):
    # Load dataset config
    classnames, concept_names = load_dataset_config(args.dataset_root, args.dataset)

    # Load backbone (only needed if we'll compute text embeddings on the fly)
    cache_present = (
        _resolve_text_embs_path(args.dataset_root, args.dataset, args.backbone, "classname") is not None
        and not args.recompute_text_embs
    )
    if cache_present and not args.use_concept_matrix:
        model, tokenizer = None, None
        print("Skipping backbone load (using cached text embeddings)")
    else:
        model, _, tokenizer = load_backbone(args.backbone, args.device)

    # Load or construct A matrix
    if args.use_concept_matrix:
        concept_matrix = _load_or_compute_text_embs(model, concept_names, args, tokenizer, kind="concept")
        A_matrix = concept_matrix.T
        print(f"Using A=Φ (concept matrix): {A_matrix.shape}")
    else:
        if args.checkpoint_path is None:
            raise ValueError("--checkpoint_path is required when --use_concept_matrix is not set")
        A_matrix = torch.load(args.checkpoint_path, map_location=args.device, weights_only=True)

    # Load test embeddings
    class_split, seen_embs, seen_targets, unseen_embs, unseen_targets = \
        load_test_embeddings(args.dataset_root, args.dataset, args.backbone)

    seen_class_ids = class_split['seen_classes']
    unseen_class_ids = class_split['unseen_classes']

    # Map global targets to seen and unseen targets
    seen_map = {global_id: local_idx for local_idx, global_id in enumerate(seen_class_ids)}
    unseen_map = {global_id: local_idx for local_idx, global_id in enumerate(unseen_class_ids)}

    # Get classname embeddings (cached or computed) and derive seen/unseen by indexing
    all_classname_embs = _load_or_compute_text_embs(model, classnames, args, tokenizer, kind="classname")
    seen_classname_embs = all_classname_embs[seen_class_ids]
    unseen_classname_embs = all_classname_embs[unseen_class_ids]

    # Get dataloaders
    seen_loader, unseen_loader = get_dataloader(
        seen_embs, seen_targets, unseen_embs, unseen_targets, args.batch_size
    )

    # Calculate quantitative results
    with torch.no_grad():
        seen_clip_acc, seen_ezpc_acc, _, _, _, _ = test(seen_loader, seen_classname_embs, seen_map, A_matrix, args.device)
        unseen_clip_acc, unseen_ezpc_acc, _, _, _, _ = test(unseen_loader, unseen_classname_embs, unseen_map, A_matrix, args.device)

        seen_gen_clip_acc, seen_gen_ezpc_acc, top1_seen, spear_seen, kendall_seen, kl_seen = test(seen_loader, all_classname_embs, None, A_matrix, args.device)
        unseen_gen_clip_acc, unseen_gen_ezpc_acc, top1_unseen, spear_unseen, kendall_unseen, kl_unseen = test(unseen_loader, all_classname_embs, None, A_matrix, args.device)

        clip_hm = 2 * seen_gen_clip_acc * unseen_gen_clip_acc / max(seen_gen_clip_acc + unseen_gen_clip_acc, 1e-8)
        ezpc_hm = 2 * seen_gen_ezpc_acc * unseen_gen_ezpc_acc / max(seen_gen_ezpc_acc + unseen_gen_ezpc_acc, 1e-8)

        n_seen = len(seen_loader.dataset)
        n_unseen = len(unseen_loader.dataset)
        n_total = n_seen + n_unseen

        top1_all = (top1_seen * n_seen + top1_unseen * n_unseen) / n_total
        spear_all = (spear_seen * n_seen + spear_unseen * n_unseen) / n_total
        kendall_all = (kendall_seen * n_seen + kendall_unseen * n_unseen) / n_total
        kl_all = (kl_seen * n_seen + kl_unseen * n_unseen) / n_total

    results = {
        "CLIP Seen Accuracy": seen_clip_acc,
        "CLIP Unseen Accuracy": unseen_clip_acc,
        "CLIP Generalized Seen Accuracy": seen_gen_clip_acc,
        "CLIP Generalized Unseen Accuracy": unseen_gen_clip_acc,
        "CLIP Generalized Harmonic Mean": clip_hm,
        "EZPC Seen Accuracy": seen_ezpc_acc,
        "EZPC Unseen Accuracy": unseen_ezpc_acc,
        "EZPC Generalized Seen Accuracy": seen_gen_ezpc_acc,
        "EZPC Generalized Unseen Accuracy": unseen_gen_ezpc_acc,
        "EZPC Generalized Harmonic Mean": ezpc_hm,
        "Fidelity Seen Top1Agree (%)": float(top1_seen),
        "Fidelity Seen Spearman": float(spear_seen),
        "Fidelity Seen Kendall (topk)": float(kendall_seen),
        "Fidelity Seen KL": float(kl_seen),
        "Fidelity Unseen Top1Agree (%)": float(top1_unseen),
        "Fidelity Unseen Spearman": float(spear_unseen),
        "Fidelity Unseen Kendall (topk)": float(kendall_unseen),
        "Fidelity Unseen KL": float(kl_unseen),
        "Fidelity All Top1Agree (%)": float(top1_all),
        "Fidelity All Spearman": float(spear_all),
        "Fidelity All Kendall (topk)": float(kendall_all),
        "Fidelity All KL": float(kl_all),
        "Fidelity Kendall topk": int(args.kendall_topk),
    }

    # Save to txt file
    os.makedirs(args.output_path, exist_ok=True)
    backbone_name = backbone_to_name(args.backbone)
    result_file = f"{args.output_path}/{args.dataset}_{backbone_name}_results.txt"
    with open(result_file, "w") as f:
        for key, value in results.items():
            f.write(f"{key}: {value}\n")

    print(f"Results saved to {result_file}")

if __name__ == "__main__":
    # Set the seed for reproducibility
    set_seed(1234)

    parser = argparse.ArgumentParser(description="Testing")
    parser.add_argument("--backbone", type=str, default="RN50",
                        help="CLIP/SigLIP backbone (e.g. RN50, ViT-B/32, ViT-L/14, siglip-so400m-patch14-384)")
    parser.add_argument("--dataset", type=str, required=True, choices=DATASET_CHOICES,
                        help="Dataset name")
    parser.add_argument("--dataset_root", type=str, required=True,
                        help="Path to the root dataset folder")
    parser.add_argument("--checkpoint_path", type=str, default=None, 
                        help="Path to the trained A matrix checkpoint")
    parser.add_argument("--use_concept_matrix", action="store_true",
                        help="Use A=Φ (raw concept embeddings as projection matrix, no training)")
    parser.add_argument("--batch_size", type=int, default=512, 
                        help="Batch size for evaluation")
    parser.add_argument("--kendall_topk", type=int, default=50, 
                        help="Top-k classes for Kendall tau approximation")
    parser.add_argument("--output_path", type=str, default="./results",
                        help="Directory to save the evaluation results")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Computation device (e.g. 'cuda', 'cpu', 'mps')")
    parser.add_argument("--recompute_text_embs", action="store_true",
                        help="Recompute text embeddings on the fly instead of loading "
                             "the cached ones next to the checkpoint")

    args = parser.parse_args()

    # Run the main function
    main(args)