| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any |
|
|
| import joblib |
| import numpy as np |
| from PIL import Image |
|
|
| from .compare_models import load_best_model_record |
| from .data_discovery import CANONICAL_LABELS |
| from .preprocessing import load_pil_image |
| from .utils import get_logger |
|
|
|
|
| LOGGER = get_logger(__name__) |
|
|
|
|
| def model_record_from_file(path: str | Path) -> dict[str, Any]: |
| path = Path(path) |
| if path.suffix == ".joblib": |
| bundle = joblib.load(path) |
| meta = bundle["metadata"] |
| return { |
| "model_name": meta["model_name"], |
| "model_type": "classical", |
| "feature_type": meta["feature_type"], |
| "model_path": str(path), |
| } |
| if path.suffix == ".pt": |
| from .dl_models import load_torch_checkpoint |
|
|
| ckpt = load_torch_checkpoint(path, map_location="cpu") |
| return { |
| "model_name": ckpt.get("model_name", ckpt["model_key"]), |
| "model_type": "deep_learning", |
| "model_key": ckpt["model_key"], |
| "family": ckpt.get("family", "cnn"), |
| "model_path": str(path), |
| } |
| raise ValueError(f"Unsupported model file: {path}") |
|
|
|
|
| def list_available_model_records(config: dict[str, Any]) -> list[dict[str, Any]]: |
| model_dir = Path(config["paths"]["model_dir"]) |
| records: list[dict[str, Any]] = [] |
| for path in sorted(model_dir.glob("*.joblib")) + sorted(model_dir.glob("*.pt")): |
| try: |
| records.append(model_record_from_file(path)) |
| except Exception as exc: |
| LOGGER.warning("Could not load model metadata for %s: %s", path, exc) |
| return records |
|
|
|
|
| class EggDamagePredictor: |
| def __init__(self, record: dict[str, Any], config: dict[str, Any]) -> None: |
| self.record = record |
| self.config = config |
| self.model_name = record["model_name"] |
| self.model_type = record["model_type"] |
| self.model_path = Path(record["model_path"]) |
| self.class_names = list(CANONICAL_LABELS) |
| self.device = None |
| if self.model_type == "classical": |
| bundle = joblib.load(self.model_path) |
| self.pipeline = bundle["pipeline"] |
| self.metadata = bundle["metadata"] |
| self.feature_type = self.metadata["feature_type"] |
| self.model = None |
| elif self.model_type == "deep_learning": |
| import torch |
|
|
| from .augmentations import build_eval_transform |
| from .dl_models import create_model, load_torch_checkpoint |
|
|
| checkpoint = load_torch_checkpoint(self.model_path, map_location="cpu") |
| self.metadata = checkpoint |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.model = create_model(checkpoint["model_key"], checkpoint.get("config", config), pretrained=False) |
| self.model.load_state_dict(checkpoint["state_dict"]) |
| self.model.to(self.device) |
| self.model.eval() |
| self.transform = build_eval_transform(checkpoint.get("config", config)) |
| self.pipeline = None |
| self.feature_type = None |
| else: |
| raise ValueError(f"Unsupported model type: {self.model_type}") |
|
|
| def predict_proba(self, image: str | Path | Image.Image | np.ndarray) -> np.ndarray: |
| pil = load_pil_image(Image.fromarray(image) if isinstance(image, np.ndarray) else image, mode="RGB") |
| if self.model_type == "classical": |
| from .classical_features import extract_single_feature |
|
|
| feature = extract_single_feature(pil, self.feature_type, self.metadata.get("config", self.config)) |
| return self.pipeline.predict_proba(feature.reshape(1, -1))[0] |
| import torch |
|
|
| assert self.model is not None and self.device is not None |
| tensor = self.transform(pil).unsqueeze(0).to(self.device) |
| with torch.no_grad(): |
| logits = self.model(tensor) |
| probs = torch.softmax(logits, dim=1).detach().cpu().numpy()[0] |
| return probs |
|
|
| def predict(self, image: str | Path | Image.Image | np.ndarray) -> dict[str, Any]: |
| probs = self.predict_proba(image) |
| pred_idx = int(np.argmax(probs)) |
| confidence = float(probs[pred_idx]) |
| return { |
| "model_name": self.model_name, |
| "model_type": self.model_type, |
| "predicted_label": self.class_names[pred_idx], |
| "predicted_index": pred_idx, |
| "confidence": confidence, |
| "probabilities": {self.class_names[i]: float(probs[i]) for i in range(len(self.class_names))}, |
| "prob_damaged": float(probs[1]), |
| "low_confidence": confidence < float(self.config["gradio"].get("low_confidence_threshold", 0.65)), |
| } |
|
|
|
|
| def load_best_predictor(config: dict[str, Any]) -> EggDamagePredictor: |
| return EggDamagePredictor(load_best_model_record(config), config) |
|
|
