| 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=[prompt.lower() for prompt in prompts], |
| padding="max_length", |
| max_length=64, |
| truncation=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() |
|
|
