Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any, Dict | |
| import torch | |
| from PIL import Image | |
| from train.dataset import load_json | |
| from .model_adapters import BaselineAdapter, GNNAdapter, InferenceResult | |
| class LoadedModel: | |
| adapter: Any | |
| label_vocab: Dict[str, int] | |
| def _detect_model_type(state: Dict[str, Any]) -> str: | |
| if "model_state" in state: | |
| return "baseline" | |
| if "gnn_state" in state and "backbone_state" in state: | |
| return "gnn" | |
| raise RuntimeError("Unrecognized checkpoint format (expected baseline or GNN keys).") | |
| def load_checkpoint(path: Path, device: torch.device | None = None) -> LoadedModel: | |
| if device is None: | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| else: | |
| device = torch.device("cpu") | |
| state: Dict[str, Any] = torch.load(path, map_location=device) | |
| model_type = _detect_model_type(state) | |
| label_vocab = state["label_vocab"] | |
| if model_type == "baseline": | |
| adapter = BaselineAdapter(state, device=device) | |
| else: | |
| adapter = GNNAdapter(state, device=device) | |
| return LoadedModel(adapter=adapter, label_vocab=label_vocab) | |
| def predict_image( | |
| image_path: Path, | |
| loaded: LoadedModel, | |
| top_k: int = 5, | |
| threshold: float = 0.5, | |
| ) -> InferenceResult: | |
| image = Image.open(image_path).convert("RGB") | |
| return loaded.adapter.predict( | |
| image=image, | |
| label_vocab=loaded.label_vocab, | |
| top_k=top_k, | |
| threshold=threshold, | |
| ) | |