src/inference.py

from __future__ import annotations

import base64
import io
import json
import os
from pathlib import Path
from typing import Dict, List, Optional

from openai import OpenAI
from PIL import Image

TORCH_IMPORT_ERROR: Optional[str] = None
try:
    import torch
    import torch.nn.functional as F
    from torchvision import models, transforms
except Exception as exc:  # pragma: no cover - defensive import guard
    torch = None
    F = None
    models = None
    transforms = None
    TORCH_IMPORT_ERROR = str(exc)

TRANSFORMERS_IMPORT_ERROR: Optional[str] = None
try:
    from transformers import CLIPModel, CLIPProcessor
except Exception as exc:  # pragma: no cover - defensive import guard
    CLIPModel = None
    CLIPProcessor = None
    TRANSFORMERS_IMPORT_ERROR = str(exc)


class CustomModelPredictor:
    def __init__(self, checkpoint_path: str = "models/custom_resnet18.pth") -> None:
        self.checkpoint_path = Path(checkpoint_path)
        self.error: Optional[str] = None
        self.model = None
        self.labels: List[str] = []
        self.image_size = 224
        self.eval_transform = None

        if torch is None or transforms is None:
            self.device = None
            self.error = (
                "Custom model unavailable because torch/torchvision is missing. "
                f"Import error: {TORCH_IMPORT_ERROR}"
            )
            return

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.eval_transform = transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ]
        )

        if self.checkpoint_path.exists():
            self._load()

    def _load(self) -> None:
        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
        self.labels = checkpoint["labels"]
        self.image_size = int(checkpoint.get("image_size", 224))

        model = models.resnet18(weights=None)
        model.fc = torch.nn.Linear(model.fc.in_features, len(self.labels))
        model.load_state_dict(checkpoint["state_dict"])
        model.to(self.device)
        model.eval()
        self.model = model

    def available(self) -> bool:
        return self.model is not None and self.error is None

    def predict(self, image: Image.Image, top_k: int = 3) -> Dict[str, object]:
        if not self.available():
            return {
                "model": "custom-transfer-learning",
                "available": False,
                "error": self.error or f"Model not found at {self.checkpoint_path}",
            }

        image = image.convert("RGB")
        tensor = self.eval_transform(image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            logits = self.model(tensor)
            probs = F.softmax(logits, dim=1).squeeze(0)

        top_probs, top_idx = torch.topk(probs, k=min(top_k, len(self.labels)))
        predictions = [
            {"label": self.labels[idx], "confidence": float(prob)}
            for prob, idx in zip(top_probs.cpu().tolist(), top_idx.cpu().tolist())
        ]

        return {
            "model": "custom-transfer-learning",
            "available": True,
            "top_prediction": predictions[0],
            "predictions": predictions,
        }


class ClipPredictor:
    def __init__(self, labels: List[str], model_name: str = "openai/clip-vit-base-patch32") -> None:
        self.labels = labels
        self.model_name = model_name
        self.error: Optional[str] = None
        self.device = torch.device("cuda" if torch is not None and torch.cuda.is_available() else "cpu") if torch is not None else None
        self.available_flag = False
        self.processor = None
        self.model = None

        if torch is None or CLIPModel is None or CLIPProcessor is None:
            self.error = (
                "CLIP unavailable because required dependencies are missing. "
                f"torch error: {TORCH_IMPORT_ERROR}; transformers error: {TRANSFORMERS_IMPORT_ERROR}"
            )
            return

        if labels:
            self._load()

    def _load(self) -> None:
        try:
            self.processor = CLIPProcessor.from_pretrained(self.model_name)
            self.model = CLIPModel.from_pretrained(self.model_name).to(self.device)
            self.model.eval()
            self.available_flag = True
        except Exception:
            self.available_flag = False

    def available(self) -> bool:
        return self.available_flag

    def predict(self, image: Image.Image, top_k: int = 3) -> Dict[str, object]:
        if not self.available():
            return {
                "model": "clip-open-source",
                "available": False,
                "error": self.error or "CLIP model could not be loaded.",
            }

        prompts = [f"a sprite or photo of a pokemon named {label}" for label in self.labels]
        inputs = self.processor(text=prompts, images=image.convert("RGB"), return_tensors="pt", padding=True)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.model(**inputs)
            logits_per_image = outputs.logits_per_image
            probs = logits_per_image.softmax(dim=1).squeeze(0)

        top_probs, top_idx = torch.topk(probs, k=min(top_k, len(self.labels)))
        predictions = [
            {"label": self.labels[idx], "confidence": float(prob)}
            for prob, idx in zip(top_probs.cpu().tolist(), top_idx.cpu().tolist())
        ]

        return {
            "model": "clip-open-source",
            "available": True,
            "top_prediction": predictions[0],
            "predictions": predictions,
        }


class OpenAIVisionPredictor:
    def __init__(self, labels: List[str], model_name: str = "gpt-4.1-mini") -> None:
        self.labels = labels
        self.model_name = model_name
        self.api_key = os.getenv("OPENAI_API_KEY", "")
        self.client: Optional[OpenAI] = None
        if self.api_key:
            self.client = OpenAI(api_key=self.api_key)

    def available(self) -> bool:
        return self.client is not None

    def predict(self, image: Image.Image) -> Dict[str, object]:
        if not self.available():
            return {
                "model": "openai-vision",
                "available": False,
                "error": "OPENAI_API_KEY is not set.",
            }

        buffered = io.BytesIO()
        image.convert("RGB").save(buffered, format="JPEG")
        b64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")

        prompt = (
            "You are an image classifier. "
            f"Choose exactly one label from this list: {', '.join(self.labels)}. "
            "Return strict JSON with keys: label, confidence, reason. "
            "label must be one of the provided labels. confidence must be in [0,1]."
        )

        response = self.client.responses.create(
            model=self.model_name,
            input=[
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": prompt},
                        {"type": "input_image", "image_url": f"data:image/jpeg;base64,{b64_image}"},
                    ],
                }
            ],
            temperature=0,
        )

        text = response.output_text.strip()
        parsed = self._safe_parse(text)
        return {
            "model": "openai-vision",
            "available": True,
            "top_prediction": {
                "label": parsed.get("label", "unknown"),
                "confidence": float(parsed.get("confidence", 0.0)),
            },
            "raw_response": parsed,
        }

    @staticmethod
    def _safe_parse(text: str) -> Dict[str, object]:
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return {"label": "unknown", "confidence": 0.0, "reason": text}