src/cv/predict.py

"""Inference wrapper around the persisted CV model."""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.runtime import configure_runtime  # noqa: E402

configure_runtime()

import torch
from PIL import Image
from torchvision import transforms

from src.config import CV_CLASSES_PATH, CV_MODEL_PATH  # noqa: E402
from src.cv.train import IMAGENET_MEAN, IMAGENET_STD, build_model  # noqa: E402


class DishClassifier:
    def __init__(self, model_path: Path | None = None) -> None:
        model_path = model_path or CV_MODEL_PATH
        if not model_path.exists():
            raise FileNotFoundError(
                f"CV model missing at {model_path}. Run 'python -m src.cv.train'."
            )
        checkpoint = torch.load(model_path, map_location="cpu")
        self.classes: list[str] = checkpoint["classes"]
        self.model_name: str = checkpoint["model_name"]
        # temperature scaling factor (set by src.cv.calibrate); 1.0 = uncalibrated
        self.temperature: float = float(checkpoint.get("temperature") or 1.0)
        self.model = build_model(self.model_name, len(self.classes))
        self.model.load_state_dict(checkpoint["state_dict"])
        self.model.eval()
        self.tf = transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
            ]
        )

    @torch.no_grad()
    def predict(self, image: Image.Image, topk: int = 3) -> list[dict[str, Any]]:
        if image.mode != "RGB":
            image = image.convert("RGB")
        x = self.tf(image).unsqueeze(0)
        logits = self.model(x)
        if self.temperature and self.temperature > 0:
            logits = logits / self.temperature
        probs = torch.softmax(logits, dim=1).squeeze(0)
        k = min(topk, len(self.classes))
        top_probs, top_idx = probs.topk(k)
        return [
            {"label": self.classes[int(i)], "confidence": float(p)}
            for p, i in zip(top_probs, top_idx)
        ]


def load_classes() -> list[str]:
    if CV_CLASSES_PATH.exists():
        return json.loads(CV_CLASSES_PATH.read_text())
    return []