test.py

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path

import open_clip
import torch
from torch.utils.data import DataLoader

PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from data import (
    CC12M_DATASET_NAMES,
    CC12MImageDataset,
    CC12MTextDataset,
    CIFAR100_DATASET_NAMES,
    CIFAR100_PROMPT_TEMPLATE,
    CIFAR100ImageDataset,
    CIFAR100TextDataset,
    IMAGENET_DATASET_NAMES,
    IMAGENET_PROMPT_TEMPLATE,
    ImageNetImageDataset,
    ImageNetTextDataset,
    MSCOCO_DATASET_NAMES,
    MSCOCOImageDataset,
    MSCOCOTextDataset,
    SUPPORTED_DATASETS,
    ReIDImageDataset,
    ReIDTextDataset,
    collate_cc12m_images,
    collate_cc12m_texts,
    collate_cifar100_images,
    collate_cifar100_texts,
    collate_imagenet_images,
    collate_imagenet_texts,
    collate_mscoco_images,
    collate_mscoco_texts,
    collate_reid_images,
    collate_reid_texts,
    load_cc12m_records,
    load_cifar100_split,
    load_imagenet_split,
    load_mscoco_records,
    load_reid_records,
    normalize_cc12m_split,
    normalize_cifar100_split,
    normalize_imagenet_split,
    normalize_mscoco_split,
)
from retrieval import encode_image_loader, encode_text_loader, retrieval_metrics
from training import get_device, seed_everything


MODEL_NAME = "ViT-B-16"
BASELINE_PRETRAINED = "laion2b_s34b_b88k"
SUPPORTED_TEST_DATASETS = tuple(
    (
        *SUPPORTED_DATASETS,
        *MSCOCO_DATASET_NAMES,
        *CC12M_DATASET_NAMES,
        *CIFAR100_DATASET_NAMES,
        *IMAGENET_DATASET_NAMES,
    )
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate text-to-image retrieval with top-k and mAP."
    )
    parser.add_argument(
        "--env-file",
        default="env/.env",
        help="Path to the .env file containing DATASET_ROOT.",
    )
    parser.add_argument("--dataset", required=True, choices=SUPPORTED_TEST_DATASETS)
    parser.add_argument(
        "--test-split",
        default="test",
        help="ReID split name. For MS COCO, the default 'test' is mapped to val2017. For CC12M, split aliases map to all. For CIFAR-100, use train or test. For ImageNet, test maps to val.",
    )
    parser.add_argument(
        "--model",
        required=True,
        help="Use 'baseline' for LAION2B ViT-B-16, or pass a checkpoint path like artifacts/model.pt.",
    )
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--top-k", type=int, nargs="+", default=[1, 5, 10])
    parser.add_argument(
        "--max-gallery-samples",
        type=int,
        default=0,
        help="Debug limit. 0 means use all, except CC12M requires a bounded sample.",
    )
    parser.add_argument(
        "--max-query-samples",
        type=int,
        default=0,
        help="Debug limit. 0 means use all captions/prompts.",
    )
    parser.add_argument(
        "--output-json",
        default="",
        help="Path to save metrics as JSON. Auto-generated under results/ if not specified.",
    )
    parser.add_argument(
        "--no-amp", action="store_true", help="Disable CUDA mixed precision."
    )
    return parser.parse_args()


def load_test_checkpoint(path: Path, device: torch.device) -> dict:
    if not path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {path}")
    checkpoint = torch.load(path, map_location=device)
    if "model_state_dict" not in checkpoint:
        raise KeyError(f"Checkpoint does not contain model_state_dict: {path}")
    return checkpoint


def is_mscoco_dataset(dataset: str) -> bool:
    return dataset.strip().lower() in MSCOCO_DATASET_NAMES


def is_cc12m_dataset(dataset: str) -> bool:
    return dataset.strip().lower() in CC12M_DATASET_NAMES


def is_cifar100_dataset(dataset: str) -> bool:
    return dataset.strip().lower() in CIFAR100_DATASET_NAMES


def is_imagenet_dataset(dataset: str) -> bool:
    return dataset.strip().lower() in IMAGENET_DATASET_NAMES


def resolve_cc12m_record_limit(
    *, max_gallery_samples: int, max_query_samples: int
) -> int:
    if max_gallery_samples < 0 or max_query_samples < 0:
        raise ValueError("--max-gallery-samples and --max-query-samples must be >= 0.")
    if max_gallery_samples:
        return max_gallery_samples
    if max_query_samples:
        return max_query_samples
    raise RuntimeError(
        "CC12M is stored as a 12.4M-row URL TSV. Pass --max-gallery-samples N "
        "(or --max-query-samples N) to evaluate a bounded sample."
    )


def make_reid_loaders(
    *,
    records,
    preprocess,
    tokenizer,
    batch_size: int,
    num_workers: int,
    max_gallery_samples: int,
    max_query_samples: int,
) -> tuple[DataLoader, DataLoader, int, int]:
    gallery_records = list(records)
    if max_gallery_samples:
        gallery_records = gallery_records[:max_gallery_samples]
    image_dataset = ReIDImageDataset(
        gallery_records,
        transform=preprocess,
    )
    text_dataset = ReIDTextDataset(
        gallery_records,
        max_samples=max_query_samples or None,
    )
    image_loader = DataLoader(
        image_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=collate_reid_images,
    )
    text_loader = DataLoader(
        text_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=lambda batch: collate_reid_texts(batch, tokenizer),
    )
    return image_loader, text_loader, len(image_dataset), len(text_dataset)


def make_cc12m_loaders(
    *,
    records,
    preprocess,
    tokenizer,
    batch_size: int,
    num_workers: int,
    max_gallery_samples: int,
    max_query_samples: int,
) -> tuple[DataLoader, DataLoader, int, int]:
    gallery_records = list(records)
    if max_gallery_samples:
        gallery_records = gallery_records[:max_gallery_samples]
    image_dataset = CC12MImageDataset(
        gallery_records,
        transform=preprocess,
    )
    text_dataset = CC12MTextDataset(
        gallery_records,
        max_samples=max_query_samples or None,
    )
    image_loader = DataLoader(
        image_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=collate_cc12m_images,
    )
    text_loader = DataLoader(
        text_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=lambda batch: collate_cc12m_texts(batch, tokenizer),
    )
    return image_loader, text_loader, len(image_dataset), len(text_dataset)


def make_cifar100_loaders(
    *,
    records,
    preprocess,
    tokenizer,
    batch_size: int,
    num_workers: int,
    max_gallery_samples: int,
    max_query_samples: int,
) -> tuple[DataLoader, DataLoader, int, int]:
    image_dataset = CIFAR100ImageDataset(
        records,
        transform=preprocess,
        max_samples=max_gallery_samples or None,
    )
    text_dataset = CIFAR100TextDataset(
        records.fine_label_names,
        class_labels=image_dataset.class_labels,
        max_samples=max_query_samples or None,
    )
    image_loader = DataLoader(
        image_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=collate_cifar100_images,
    )
    text_loader = DataLoader(
        text_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=lambda batch: collate_cifar100_texts(batch, tokenizer),
    )
    return image_loader, text_loader, len(image_dataset), len(text_dataset)


def make_imagenet_loaders(
    *,
    records,
    preprocess,
    tokenizer,
    batch_size: int,
    num_workers: int,
    max_gallery_samples: int,
    max_query_samples: int,
) -> tuple[DataLoader, DataLoader, int, int]:
    image_dataset = ImageNetImageDataset(
        records,
        transform=preprocess,
        max_samples=max_gallery_samples or None,
    )
    text_dataset = ImageNetTextDataset(
        records.class_names,
        class_labels=image_dataset.class_labels,
        max_samples=max_query_samples or None,
    )
    image_loader = DataLoader(
        image_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=collate_imagenet_images,
    )
    text_loader = DataLoader(
        text_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=lambda batch: collate_imagenet_texts(batch, tokenizer),
    )
    return image_loader, text_loader, len(image_dataset), len(text_dataset)


def make_mscoco_loaders(
    *,
    records,
    preprocess,
    tokenizer,
    batch_size: int,
    num_workers: int,
    max_gallery_samples: int,
    max_query_samples: int,
) -> tuple[DataLoader, DataLoader, int, int]:
    gallery_records = list(records)
    if max_gallery_samples:
        gallery_records = gallery_records[:max_gallery_samples]
    image_dataset = MSCOCOImageDataset(
        gallery_records,
        transform=preprocess,
    )
    text_dataset = MSCOCOTextDataset(
        gallery_records,
        max_samples=max_query_samples or None,
    )
    image_loader = DataLoader(
        image_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=collate_mscoco_images,
    )
    text_loader = DataLoader(
        text_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        collate_fn=lambda batch: collate_mscoco_texts(batch, tokenizer),
    )
    return image_loader, text_loader, len(image_dataset), len(text_dataset)


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    device = get_device(args.device)
    dataset_name = args.dataset.strip().lower()
    use_mscoco = is_mscoco_dataset(dataset_name)
    use_cc12m = is_cc12m_dataset(dataset_name)
    use_cifar100 = is_cifar100_dataset(dataset_name)
    use_imagenet = is_imagenet_dataset(dataset_name)
    eval_split = args.test_split
    if use_mscoco and eval_split == "test":
        eval_split = "val2017"
    is_baseline = args.model == "baseline"
    checkpoint_path = None if is_baseline else Path(args.model)
    init_pretrained = BASELINE_PRETRAINED if is_baseline else ""

    if use_mscoco:
        eval_split = normalize_mscoco_split(eval_split)
        test_records = load_mscoco_records(env_file=args.env_file, split=eval_split)
    elif use_cc12m:
        eval_split = normalize_cc12m_split(eval_split)
        cc12m_record_limit = resolve_cc12m_record_limit(
            max_gallery_samples=args.max_gallery_samples,
            max_query_samples=args.max_query_samples,
        )
        test_records = load_cc12m_records(
            env_file=args.env_file,
            split=eval_split,
            max_records=cc12m_record_limit,
        )
    elif use_cifar100:
        eval_split = normalize_cifar100_split(eval_split)
        test_records = load_cifar100_split(env_file=args.env_file, split=eval_split)
    elif use_imagenet:
        eval_split = normalize_imagenet_split(eval_split)
        test_records = load_imagenet_split(env_file=args.env_file, split=eval_split)
    else:
        records = load_reid_records(env_file=args.env_file, dataset=args.dataset)
        test_records = [record for record in records if record.split == eval_split]
    if use_cifar100 and not test_records.fine_labels:
        raise RuntimeError(f"No records found for test split {eval_split!r}.")
    if use_imagenet and not test_records.records:
        raise RuntimeError(f"No records found for test split {eval_split!r}.")
    if not use_cifar100 and not use_imagenet and not test_records:
        raise RuntimeError(f"No records found for test split {eval_split!r}.")

    model, _, preprocess = open_clip.create_model_and_transforms(
        MODEL_NAME,
        pretrained=init_pretrained,
        device=device,
    )
    saved_args: dict = {}
    if checkpoint_path is not None:
        checkpoint = load_test_checkpoint(checkpoint_path, device)
        saved_args = checkpoint.get("args", {})
        if not isinstance(saved_args, dict):
            saved_args = {}
        checkpoint_model_name = saved_args.get("model_name")
        if checkpoint_model_name and checkpoint_model_name != MODEL_NAME:
            raise ValueError(
                f"Checkpoint was trained with {checkpoint_model_name!r}, but test.py is fixed to {MODEL_NAME!r}."
            )
        model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    tokenizer = open_clip.get_tokenizer(MODEL_NAME)
    use_amp = device.type == "cuda" and not args.no_amp

    if use_mscoco:
        loader_factory = make_mscoco_loaders
        query_unit = "captions"
        domain = "general-image-text"
    elif use_cc12m:
        loader_factory = make_cc12m_loaders
        query_unit = "captions"
        domain = "web-image-text"
    elif use_cifar100:
        loader_factory = make_cifar100_loaders
        query_unit = "class prompts"
        domain = "general-classification"
    elif use_imagenet:
        loader_factory = make_imagenet_loaders
        query_unit = "class prompts"
        domain = "general-classification"
    else:
        loader_factory = make_reid_loaders
        query_unit = "captions"
        domain = "person-reid"
    image_loader, text_loader, gallery_count, query_count = loader_factory(
        records=test_records,
        preprocess=preprocess,
        tokenizer=tokenizer,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        max_gallery_samples=args.max_gallery_samples,
        max_query_samples=args.max_query_samples,
    )
    print(
        f"Evaluating dataset={dataset_name} model={args.model} split={eval_split} "
        f"queries={query_count:,} {query_unit}, gallery={gallery_count:,} images on {device}"
    )

    image_features, image_ids = encode_image_loader(
        model, image_loader, device=device, use_amp=use_amp
    )
    text_features, text_ids = encode_text_loader(
        model, text_loader, device=device, use_amp=use_amp
    )
    if use_cc12m:
        gallery_id_set = set(int(target_id) for target_id in image_ids.tolist())
        keep_query_mask = torch.tensor(
            [int(target_id) in gallery_id_set for target_id in text_ids.tolist()],
            dtype=torch.bool,
        )
        if not bool(keep_query_mask.any()):
            raise RuntimeError(
                "No CC12M text queries have a successfully loaded gallery image."
            )
        dropped_queries = int(text_ids.shape[0] - keep_query_mask.sum().item())
        if image_ids.shape[0] != gallery_count or dropped_queries:
            print(
                f"CC12M usable gallery={image_ids.shape[0]:,}/{gallery_count:,}; "
                f"dropped_queries={dropped_queries:,}"
            )
        text_features = text_features[keep_query_mask]
        text_ids = text_ids[keep_query_mask]
    metrics = retrieval_metrics(
        query_features=text_features,
        gallery_features=image_features,
        query_ids=text_ids,
        gallery_ids=image_ids,
        top_k=args.top_k,
    )

    print("text-to-image retrieval")
    for key in [f"top{k}" for k in args.top_k] + ["mAP"]:
        print(f"{key}: {metrics[key]:.4f}")

    output = {
        "model": args.model,
        "dataset": dataset_name,
        "domain": domain,
        "checkpoint": "" if checkpoint_path is None else str(checkpoint_path),
        "model_name": MODEL_NAME,
        "baseline_pretrained": BASELINE_PRETRAINED if is_baseline else "",
        "checkpoint_pretrained": saved_args.get("pretrained", ""),
        "init_pretrained": init_pretrained,
        "split": eval_split,
        "prompt_template": (
            CIFAR100_PROMPT_TEMPLATE
            if use_cifar100
            else IMAGENET_PROMPT_TEMPLATE
            if use_imagenet
            else ""
        ),
        "queries": query_count,
        "gallery": gallery_count,
        "encoded_queries": int(text_features.shape[0]),
        "encoded_gallery": int(image_features.shape[0]),
        "metrics": metrics,
    }
    if args.output_json:
        output_path = Path(args.output_json)
    else:
        model_tag = "baseline" if is_baseline else Path(args.model).stem
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = Path("results") / f"{dataset_name}_{model_tag}_{timestamp}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(output, indent=2, ensure_ascii=False), encoding="utf-8"
    )
    print(f"Results saved to {output_path}")


if __name__ == "__main__":
    main()