Spaces:
Sleeping
Sleeping
| # filename: src/inference/model_loader.py | |
| # purpose: Load all inference models at FastAPI startup. | |
| # Checks local models/ first, falls back to HF Hub with 120s timeout. | |
| # Returns (models_dict, hashes_dict) — hashes stored separately from nn.Module. | |
| from __future__ import annotations | |
| import hashlib | |
| import logging | |
| from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout | |
| from pathlib import Path | |
| from typing import Any, Optional | |
| import torch | |
| import torch.nn as nn | |
| import config as cfg | |
| from src.models.base_classifier import load_model | |
| from src.models.model_factory import get_model | |
| from src.utils.helpers import download_model_from_hub | |
| logger = logging.getLogger(__name__) | |
| CNN_WEIGHT_FILES: dict[str, str] = { | |
| "resnet50": "resnet50_best.pt", | |
| "mobilenet": "mobilenet_best.pt", | |
| } | |
| YOLO_WEIGHT_PATH: Path = cfg.ARTIFACTS_DIR / "detection" / "yolov8_smartvision.pt" | |
| HF_DOWNLOAD_TIMEOUT: int = 120 | |
| def load_all_models(device: str = "cpu") -> tuple[dict[str, Any], dict[str, str]]: | |
| """Load all inference models. Returns (models_dict, hashes_dict). | |
| models_dict keys: "resnet50", "mobilenet" (if FAST_MODE=False), "yolo" | |
| hashes_dict keys: same as models_dict (excludes "yolo") | |
| Gracefully skips any model that fails to load — caller checks which are present. | |
| FAST_MODE=True (default): ResNet50 + YOLO only. | |
| FAST_MODE=False: ResNet50 + MobileNet + YOLO. | |
| """ | |
| models: dict[str, Any] = {} | |
| hashes: dict[str, str] = {} | |
| cnn_names = ["resnet50"] | |
| if not cfg.FAST_MODE: | |
| cnn_names.append("mobilenet") | |
| for name in cnn_names: | |
| model, model_hash = _load_cnn(name, device) | |
| if model is not None: | |
| models[name] = model | |
| hashes[name] = model_hash | |
| logger.info(f"Loaded CNN: {name} (hash={model_hash})") | |
| else: | |
| logger.warning(f"Skipping {name} — weight file not available.") | |
| yolo = _load_yolo() | |
| if yolo is not None: | |
| models["yolo"] = yolo | |
| logger.info("Loaded YOLO detector.") | |
| else: | |
| logger.warning("Skipping YOLO — weight file not available.") | |
| _warm_up(models, device) | |
| return models, hashes | |
| def _load_cnn(name: str, device: str) -> tuple[Optional[nn.Module], str]: | |
| """Build architecture + load weights. Returns (model, hash_8chars) or (None, '').""" | |
| weight_path = _resolve_weight_path(name) | |
| if weight_path is None: | |
| return None, "" | |
| try: | |
| model_hash = _file_hash(weight_path) | |
| model = get_model(name, num_classes=cfg.NUM_CLASSES, pretrained=False) | |
| model = load_model(model, weight_path, device=device) | |
| model.eval() | |
| return model, model_hash | |
| except Exception as e: | |
| logger.error(f"Failed to load CNN {name}: {e}") | |
| return None, "" | |
| def _file_hash(path: Path) -> str: | |
| """MD5 via 8MB chunks — avoids reading full file (100MB+) into memory.""" | |
| hasher = hashlib.md5() | |
| with open(path, "rb") as f: | |
| while chunk := f.read(8 * 1024 * 1024): | |
| hasher.update(chunk) | |
| return hasher.hexdigest()[:8] | |
| def _load_yolo() -> Optional[Any]: | |
| """Load YOLOv8n from local weight file. | |
| NOTE: YOLO .pt files require weights_only=False because they embed custom | |
| ultralytics classes (DetectionModel, etc.). This is safe here because the | |
| file was generated by our own training run in Section 6. | |
| """ | |
| if not YOLO_WEIGHT_PATH.exists(): | |
| logger.error(f"YOLO weight not found: {YOLO_WEIGHT_PATH}") | |
| return None | |
| try: | |
| from ultralytics import YOLO | |
| return YOLO(str(YOLO_WEIGHT_PATH)) | |
| except Exception as e: | |
| logger.error(f"Failed to load YOLO: {e}") | |
| return None | |
| def _warm_up(models: dict[str, Any], device: str) -> None: | |
| """3 forward passes per CNN to stabilize CPU JIT and kernel caching. | |
| YOLO excluded — ultralytics handles its own internal warmup. | |
| """ | |
| cnn_models = {k: v for k, v in models.items() if k != "yolo"} | |
| if not cnn_models: | |
| return | |
| dummy = torch.randn(1, 3, cfg.IMAGE_SIZE, cfg.IMAGE_SIZE).to(device) | |
| with torch.no_grad(): | |
| for name, model in cnn_models.items(): | |
| for _ in range(3): | |
| _ = model(dummy) | |
| logger.info(f"Warmup complete (3 passes): {list(cnn_models.keys())}") | |
| def _resolve_weight_path(name: str) -> Optional[Path]: | |
| """Check local models/ dir first, then try HF Hub download.""" | |
| local_path = cfg.MODELS_DIR / CNN_WEIGHT_FILES[name] | |
| if local_path.exists(): | |
| logger.info(f"Found local weights: {local_path}") | |
| return local_path | |
| logger.info(f"{local_path.name} not found locally. Attempting HF Hub download...") | |
| return _download_with_timeout(name) | |
| def _download_with_timeout(name: str) -> Optional[Path]: | |
| """Download weight from HF Hub with HF_DOWNLOAD_TIMEOUT second limit. | |
| Returns local Path on success, None on timeout or any error. | |
| """ | |
| if ( | |
| not cfg.HF_TOKEN | |
| or not cfg.HF_REPO_ID | |
| or cfg.HF_REPO_ID == "your-username/smartvision-models" | |
| ): | |
| logger.warning( | |
| f"HF_TOKEN or HF_REPO_ID not configured. Cannot download {name}. " | |
| "Set HF_TOKEN and HF_REPO_ID in .env to enable automatic download." | |
| ) | |
| return None | |
| filename = CNN_WEIGHT_FILES[name] | |
| def _do_download() -> Path: | |
| return download_model_from_hub( | |
| filename=filename, | |
| local_dir=cfg.MODELS_DIR, | |
| repo_id=cfg.HF_REPO_ID, | |
| token=cfg.HF_TOKEN, | |
| ) | |
| with ThreadPoolExecutor(max_workers=1) as executor: | |
| future = executor.submit(_do_download) | |
| try: | |
| return future.result(timeout=HF_DOWNLOAD_TIMEOUT) | |
| except FuturesTimeout: | |
| logger.error( | |
| f"HF Hub download for {name} timed out after {HF_DOWNLOAD_TIMEOUT}s. " | |
| "Model will not be loaded." | |
| ) | |
| return None | |
| except Exception as e: | |
| logger.error(f"HF Hub download failed for {name}: {e}") | |
| return None | |