Spaces:
Sleeping
Sleeping
File size: 1,582 Bytes
419fc11 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | 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
@dataclass
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,
)
|