cifar_classifier.py

"""Self-contained CIFAR-10 classifier loader for the deployed models.

No dependency on the training repo — only torch, torchvision, timm and
huggingface_hub. Works with a local model folder or a Hugging Face Hub repo id.

Each model folder contains:
  - config.json            model + preprocessing metadata
  - model.safetensors      weights (or pytorch_model.bin as fallback)
  - README.md              model card

Example:
    from cifar_classifier import CIFAR10Classifier

    clf = CIFAR10Classifier.from_pretrained("YOUR_USER/convnextv2-huge-cifar10-upsample")
    label, probs = clf.predict("cat.jpg")
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import timm
from PIL import Image
from torchvision import transforms

CIFAR10_CLASSES = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck",
]

_WEIGHT_FILES = ("model.safetensors", "pytorch_model.bin")


def _replace_conv(conv: nn.Conv2d, stride: int = 1) -> nn.Conv2d:
    """3x3 stride-1 stem conv used for native 32x32 CIFAR input.

    Mirrors ``adapt_stem_for_cifar`` in the training repo
    (src/models/registry.py). Weights are loaded from the checkpoint, so the
    averaged-kernel initialisation used at training time is irrelevant here —
    only the shapes must match.
    """
    return nn.Conv2d(
        conv.in_channels,
        conv.out_channels,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=conv.bias is not None,
    )


def adapt_stem_for_cifar(model: nn.Module, model_name: str) -> None:
    """Adapt timm stems for native 32x32 CIFAR input (must match training)."""
    if hasattr(model, "patch_embed") and hasattr(model.patch_embed, "proj"):
        model.patch_embed.proj = _replace_conv(model.patch_embed.proj)
        return
    if hasattr(model, "stem"):
        stem = model.stem
        if hasattr(stem, "conv") and isinstance(stem.conv, nn.Conv2d):
            stem.conv = _replace_conv(stem.conv)
        elif isinstance(stem, nn.Sequential):
            for module in stem.modules():
                if isinstance(module, nn.Conv2d):
                    module.kernel_size = (3, 3)
                    module.stride = (1, 1)
                    module.padding = (1, 1)
                    break
    if model_name.startswith("convnext") and hasattr(model, "downsample_layers"):
        first = model.downsample_layers[0]
        if isinstance(first, nn.Sequential) and isinstance(first[0], nn.Conv2d):
            first[0] = _replace_conv(first[0])
    if model_name.startswith("tf_efficientnet") and hasattr(model, "conv_stem"):
        model.conv_stem = _replace_conv(model.conv_stem)


def build_model(config: dict[str, Any]) -> nn.Module:
    """Build the timm model exactly as configured at training time."""
    model = timm.create_model(
        config["timm_name"],
        pretrained=False,
        num_classes=config.get("num_classes", 10),
    )
    if config.get("input_mode", "upsample") == "native":
        adapt_stem_for_cifar(model, config["timm_name"])
    return model


def build_eval_transform(config: dict[str, Any]) -> transforms.Compose:
    """Deterministic preprocessing matching the training eval pipeline."""
    ops: list[Any] = []
    size = config.get("image_size", 32)
    if config.get("input_mode", "upsample") == "upsample":
        ops.append(transforms.Resize((size, size)))
    else:
        # Native models were trained on 32x32; downscale larger inputs.
        ops.append(transforms.Resize((size, size)))
    ops.append(transforms.ToTensor())
    ops.append(transforms.Normalize(mean=config["mean"], std=config["std"]))
    return transforms.Compose(ops)


def _load_state_dict(path: Path) -> dict[str, torch.Tensor]:
    if path.suffix == ".safetensors":
        from safetensors.torch import load_file

        return load_file(str(path))
    return torch.load(path, map_location="cpu", weights_only=True)


class CIFAR10Classifier:
    """Inference wrapper around a deployed CIFAR-10 model folder."""

    def __init__(self, model: nn.Module, config: dict[str, Any], device: str | torch.device = "cpu"):
        self.config = config
        self.classes: list[str] = config.get("classes", CIFAR10_CLASSES)
        self.device = torch.device(device)
        self.model = model.to(self.device).eval()
        self.transform = build_eval_transform(config)

    # -- loading ---------------------------------------------------------

    @classmethod
    def from_pretrained(
        cls,
        repo_id_or_path: str | os.PathLike[str],
        device: str | torch.device | None = None,
        revision: str | None = None,
    ) -> "CIFAR10Classifier":
        """Load from a local folder or a Hugging Face Hub repo id."""
        folder = cls._resolve_folder(repo_id_or_path, revision)
        with open(folder / "config.json", encoding="utf-8") as fh:
            config = json.load(fh)
        model = build_model(config)
        for name in _WEIGHT_FILES:
            weight_path = folder / name
            if weight_path.exists():
                model.load_state_dict(_load_state_dict(weight_path))
                break
        else:
            raise FileNotFoundError(f"No weight file ({', '.join(_WEIGHT_FILES)}) found in {folder}")
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        return cls(model, config, device)

    @staticmethod
    def _resolve_folder(repo_id_or_path: str | os.PathLike[str], revision: str | None) -> Path:
        path = Path(repo_id_or_path)
        if path.is_dir():
            return path
        from huggingface_hub import snapshot_download

        return Path(
            snapshot_download(
                repo_id=str(repo_id_or_path),
                revision=revision,
                allow_patterns=["config.json", "README.md", *_WEIGHT_FILES],
            )
        )

    # -- inference ---------------------------------------------------------

    def _preprocess(self, image: str | os.PathLike[str] | Image.Image) -> torch.Tensor:
        if not isinstance(image, Image.Image):
            image = Image.open(image)
        image = image.convert("RGB")
        return self.transform(image)

    @torch.no_grad()
    def predict(self, image: str | os.PathLike[str] | Image.Image) -> tuple[str, dict[str, float]]:
        """Return (top label, {class: probability}) for one image."""
        tensor = self._preprocess(image).unsqueeze(0).to(self.device)
        probs = torch.softmax(self.model(tensor), dim=-1)[0].cpu()
        top = int(probs.argmax())
        return self.classes[top], {cls: float(probs[i]) for i, cls in enumerate(self.classes)}

    @torch.no_grad()
    def predict_batch(self, images: list[str | os.PathLike[str] | Image.Image]) -> list[str]:
        """Return the top label for each image."""
        batch = torch.stack([self._preprocess(img) for img in images]).to(self.device)
        indices = self.model(batch).argmax(dim=-1).cpu().tolist()
        return [self.classes[i] for i in indices]


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Run CIFAR-10 inference with a deployed model.")
    parser.add_argument("model", help="Local model folder or Hub repo id (e.g. user/convnextv2-huge-cifar10-upsample)")
    parser.add_argument("image", help="Path to an image file")
    parser.add_argument("--device", default=None, help="cpu | cuda (default: auto)")
    args = parser.parse_args()

    clf = CIFAR10Classifier.from_pretrained(args.model, device=args.device)
    label, probs = clf.predict(args.image)
    print(f"Prediction: {label}")
    for cls_name, p in sorted(probs.items(), key=lambda kv: kv[1], reverse=True):
        print(f"  {cls_name:<12s} {p:.4f}")