import argparse import json import math import time import traceback from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import torch import torch.nn.functional as F from PIL import Image from torch.utils.data import DataLoader, Dataset from tqdm import tqdm IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} DEFAULT_MODELS = [ "cropvlm", "openai_clip_vit_b32", "bioclip", "bioclip2", "biotrove_clip", "remoteclip", "siglip2", ] class ImageFolderPaths(Dataset): def __init__(self, root: str): self.root = Path(root) self.classes = sorted([p.name for p in self.root.iterdir() if p.is_dir()]) self.class_to_idx = {name: idx for idx, name in enumerate(self.classes)} self.samples: List[Tuple[Path, int]] = [] for class_name in self.classes: for path in sorted((self.root / class_name).iterdir()): if path.is_file() and path.suffix.lower() in IMAGE_EXTS: self.samples.append((path, self.class_to_idx[class_name])) def __len__(self) -> int: return len(self.samples) def __getitem__(self, idx: int): path, label = self.samples[idx] return Image.open(path).convert("RGB"), label, str(path) def pil_collate(batch): images, labels, paths = zip(*batch) return list(images), torch.tensor(labels, dtype=torch.long), list(paths) def display_name(class_name: str) -> str: return class_name.replace("_", " ") def normalize(features: torch.Tensor) -> torch.Tensor: if isinstance(features, (tuple, list)): features = features[0] return F.normalize(features.float(), dim=-1) class Adapter: name = "" family = "" checkpoint: Optional[str] = None load_message: Optional[str] = None def encode_text(self, prompts: Sequence[str]) -> torch.Tensor: raise NotImplementedError def encode_images(self, images: Sequence[Image.Image]) -> torch.Tensor: raise NotImplementedError class OpenAIClipAdapter(Adapter): def __init__(self, device: torch.device, checkpoint: Optional[str] = None): import clip self.name = "CropVLM" if checkpoint else "OpenAI CLIP ViT-B/32" self.family = "openai_clip" self.device = device self.clip = clip self.model, self.preprocess = clip.load("ViT-B/32", device=str(device)) if checkpoint: checkpoint_path = Path(checkpoint) if not checkpoint_path.exists(): raise FileNotFoundError(f"CropVLM checkpoint not found: {checkpoint_path}") ckpt = torch.load(checkpoint_path, map_location=device) state = ckpt.get("model_state_dict", ckpt.get("state_dict", ckpt)) self.model.load_state_dict(state) self.checkpoint = str(checkpoint_path) self.model.eval() def encode_text(self, prompts: Sequence[str]) -> torch.Tensor: tokens = self.clip.tokenize(list(prompts), truncate=True).to(self.device) with torch.no_grad(): return normalize(self.model.encode_text(tokens)) def encode_images(self, images: Sequence[Image.Image]) -> torch.Tensor: batch = torch.stack([self.preprocess(image) for image in images]).to(self.device) with torch.no_grad(): return normalize(self.model.encode_image(batch)) class OpenClipAdapter(Adapter): def __init__( self, model_name: str, pretrained: Optional[str], device: torch.device, hf_checkpoint: Optional[Tuple[str, str]] = None, ): import open_clip self.name = model_name self.family = "open_clip" self.device = device self.model_name = model_name self.pretrained = pretrained self.open_clip = open_clip if hf_checkpoint: from huggingface_hub import hf_hub_download repo, filename = hf_checkpoint checkpoint = hf_hub_download(repo, filename) self.model, _, self.preprocess = open_clip.create_model_and_transforms(model_name, pretrained=None) ckpt = torch.load(checkpoint, map_location="cpu") state = ckpt.get("state_dict", ckpt.get("model_state_dict", ckpt)) if isinstance(ckpt, dict) else ckpt if any(key.startswith("module.") for key in state): state = {key.removeprefix("module."): value for key, value in state.items()} self.load_message = str(self.model.load_state_dict(state, strict=False)) self.checkpoint = checkpoint else: self.model, _, self.preprocess = open_clip.create_model_and_transforms( model_name, pretrained=pretrained, ) self.tokenizer = open_clip.get_tokenizer(model_name) self.model.to(device).eval() def encode_text(self, prompts: Sequence[str]) -> torch.Tensor: tokens = self.tokenizer(list(prompts)).to(self.device) with torch.no_grad(): return normalize(self.model.encode_text(tokens)) def encode_images(self, images: Sequence[Image.Image]) -> torch.Tensor: batch = torch.stack([self.preprocess(image) for image in images]).to(self.device) with torch.no_grad(): return normalize(self.model.encode_image(batch)) class Siglip2Adapter(Adapter): def __init__(self, device: torch.device): from transformers import AutoModel, AutoProcessor self.name = "google/siglip2-base-patch16-224" self.family = "transformers_siglip2" self.device = device self.processor = AutoProcessor.from_pretrained(self.name) self.model = AutoModel.from_pretrained(self.name).to(device).eval() def encode_text(self, prompts: Sequence[str]) -> torch.Tensor: inputs = self.processor(text=list(prompts), padding=True, return_tensors="pt").to(self.device) with torch.no_grad(): if hasattr(self.model, "get_text_features"): features = self.model.get_text_features(**inputs) else: features = self.model(**inputs).text_embeds return normalize(features) def encode_images(self, images: Sequence[Image.Image]) -> torch.Tensor: inputs = self.processor(images=list(images), return_tensors="pt").to(self.device) with torch.no_grad(): if hasattr(self.model, "get_image_features"): features = self.model.get_image_features(**inputs) else: features = self.model(**inputs).image_embeds return normalize(features) def build_adapter(model_key: str, device: torch.device, cropvlm_checkpoint: str) -> Adapter: if model_key == "cropvlm": return OpenAIClipAdapter(device, checkpoint=cropvlm_checkpoint) if model_key == "openai_clip_vit_b32": return OpenAIClipAdapter(device) if model_key == "bioclip": return OpenClipAdapter("hf-hub:imageomics/bioclip", None, device) if model_key == "bioclip2": return OpenClipAdapter("hf-hub:imageomics/bioclip-2", None, device) if model_key == "biotrove_clip": return OpenClipAdapter( "ViT-B-16", None, device, hf_checkpoint=("BGLab/BioTrove-CLIP", "biotroveclip-vit-b-16-from-bioclip-epoch-8.pt"), ) if model_key == "remoteclip": return OpenClipAdapter( "ViT-B-32", None, device, hf_checkpoint=("chendelong/RemoteCLIP", "RemoteCLIP-ViT-B-32.pt"), ) if model_key == "siglip2": return Siglip2Adapter(device) raise KeyError( f"Unknown model '{model_key}'. Supported models: {', '.join(DEFAULT_MODELS)}. " "TULIP, EVA-CLIP, and LongCLIP are intentionally excluded." ) def per_class_stats(per_class: Dict[str, Dict[str, Any]]) -> Dict[str, Any]: values = [item["accuracy"] for item in per_class.values() if item.get("accuracy") is not None] if not values: return { "per_class_accuracy_mean": None, "per_class_accuracy_std": None, "per_class_accuracy_std_population": None, "num_classes_with_accuracy": 0, } mean = sum(values) / len(values) sample_std = math.sqrt(sum((x - mean) ** 2 for x in values) / (len(values) - 1)) if len(values) > 1 else 0.0 population_std = math.sqrt(sum((x - mean) ** 2 for x in values) / len(values)) return { "per_class_accuracy_mean": mean, "per_class_accuracy_std": sample_std, "per_class_accuracy_std_population": population_std, "num_classes_with_accuracy": len(values), } def evaluate_model(args: argparse.Namespace, dataset: ImageFolderPaths, model_key: str) -> Dict[str, Any]: started_at = time.time() device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) prompts = [args.prompt_template.format(display_name(class_name)) for class_name in dataset.classes] result: Dict[str, Any] = { "model_key": model_key, "dataset": str(dataset.root), "num_images": len(dataset), "num_classes": len(dataset.classes), "classes": dataset.classes, "class_prompts": dict(zip(dataset.classes, prompts)), "prompt_template": args.prompt_template, "device": str(device), "status": "started", "started_at_unix": started_at, } try: adapter = build_adapter(model_key, device, args.cropvlm_checkpoint) result["model_name"] = adapter.name result["family"] = adapter.family result["checkpoint"] = adapter.checkpoint result["load_message"] = adapter.load_message text_features = adapter.encode_text(prompts).to(device) loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=pil_collate, ) class_total = [0 for _ in dataset.classes] class_correct = [0 for _ in dataset.classes] confusion = [[0 for _ in dataset.classes] for _ in dataset.classes] predictions: List[Dict[str, Any]] = [] correct = 0 for images, labels, paths in tqdm(loader, desc=model_key): image_features = adapter.encode_images(images) logits = image_features @ text_features.T pred = logits.argmax(dim=-1).detach().cpu() scores = logits.max(dim=-1).values.detach().cpu() for true_idx, pred_idx, score, path in zip(labels.tolist(), pred.tolist(), scores.tolist(), paths): class_total[true_idx] += 1 class_correct[true_idx] += int(true_idx == pred_idx) confusion[true_idx][pred_idx] += 1 correct += int(true_idx == pred_idx) if args.save_predictions: predictions.append( { "path": path, "true_class": dataset.classes[true_idx], "pred_class": dataset.classes[pred_idx], "correct": true_idx == pred_idx, "score": float(score), } ) per_class = {} for idx, class_name in enumerate(dataset.classes): total = class_total[idx] per_class[class_name] = { "correct": class_correct[idx], "total": total, "accuracy": class_correct[idx] / total if total else None, } result.update( { "status": "ok", "accuracy": correct / len(dataset) if len(dataset) else None, "correct": correct, "per_class": per_class, "confusion_matrix": confusion, "predictions": predictions if args.save_predictions else None, } ) result.update(per_class_stats(per_class)) except Exception as exc: result.update( { "status": "failed", "error_type": type(exc).__name__, "error": str(exc), "traceback": traceback.format_exc(), } ) result["elapsed_seconds"] = time.time() - started_at return result def write_json(path: Path, data: Dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w") as f: json.dump(data, f, indent=2) def main(): parser = argparse.ArgumentParser() parser.add_argument("--dataset", required=True, help="ImageFolder-style dataset root.") parser.add_argument("--output", default="outputs/zero_shot_results.json") parser.add_argument("--cropvlm-checkpoint", default="models/CropVLM.pth") parser.add_argument("--models", nargs="+", default=DEFAULT_MODELS) parser.add_argument("--device", default=None) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--num-workers", type=int, default=2) parser.add_argument("--prompt-template", default="{}") parser.add_argument("--save-predictions", action="store_true") args = parser.parse_args() excluded = {"tulip", "eva_clip", "eva_clip_official", "longclip"} requested = [model for model in args.models if model not in excluded] skipped = [model for model in args.models if model in excluded] dataset = ImageFolderPaths(args.dataset) results = [evaluate_model(args, dataset, model_key) for model_key in requested] ok = [result for result in results if result.get("status") == "ok"] failed = [result for result in results if result.get("status") != "ok"] summary = { "created_at": datetime.now(timezone.utc).isoformat(), "dataset": str(dataset.root), "num_images": len(dataset), "num_classes": len(dataset.classes), "classes": dataset.classes, "requested_models": args.models, "evaluated_models": requested, "skipped_models": skipped, "num_models": len(results), "num_successful": len(ok), "num_failed": len(failed), "models": { result["model_key"]: { "status": result.get("status"), "accuracy": result.get("accuracy"), "correct": result.get("correct"), "num_images": result.get("num_images"), "per_class_accuracy_mean": result.get("per_class_accuracy_mean"), "per_class_accuracy_std": result.get("per_class_accuracy_std"), "per_class_accuracy_std_population": result.get("per_class_accuracy_std_population"), "num_classes_with_accuracy": result.get("num_classes_with_accuracy"), "elapsed_seconds": result.get("elapsed_seconds"), "error": result.get("error"), } for result in results }, "model_results": {result["model_key"]: result for result in results}, "results": results, } write_json(Path(args.output), summary) print(Path(args.output)) if __name__ == "__main__": main()