from __future__ import annotations import hashlib import re from pathlib import Path from app.core.config import Settings, get_settings from app.core.errors import AppError from app.schemas.detection import ModelAssetListResponse, ModelAssetRead class ModelAssetCatalogService: SUPPORTED_SUFFIXES = { ".pt": "ultralytics/pytorch", ".onnx": "onnx", ".engine": "tensorrt", } @staticmethod def list_assets(settings: Settings | None = None) -> ModelAssetListResponse: resolved_settings = settings or get_settings() model_directory = ModelAssetCatalogService._model_directory(resolved_settings) active_model_path = ModelAssetCatalogService._resolved_file_path(resolved_settings.yolo_model_path) if not model_directory.exists() or not model_directory.is_dir(): return ModelAssetListResponse(items=[], total=0, model_directory=str(model_directory)) candidate_paths = [ path for path in sorted(model_directory.iterdir(), key=lambda item: item.name.lower()) if path.is_file() and path.suffix.lower() in ModelAssetCatalogService.SUPPORTED_SUFFIXES ] if active_model_path is not None: candidate_paths = [path for path in candidate_paths if path.resolve() == active_model_path] items = [ ModelAssetCatalogService._asset_from_file(path, active_model_path=active_model_path) for path in candidate_paths ] return ModelAssetListResponse(items=items, total=len(items), model_directory=str(model_directory)) @staticmethod def resolve_asset(model_asset_id: str, settings: Settings | None = None) -> ModelAssetRead: normalized = model_asset_id.strip() for asset in ModelAssetCatalogService.list_assets(settings=settings).items: if asset.model_asset_id == normalized: return asset raise AppError( code="DETECTION_MODEL_ASSET_NOT_FOUND", message="Selected local model asset was not found in the configured model directory", details={"model_asset_id": normalized}, status_code=404, ) @staticmethod def settings_for_asset(settings: Settings, asset: ModelAssetRead) -> Settings: return settings.model_copy(update={"yolo_model_path": asset.model_path}) @staticmethod def _model_directory(settings: Settings) -> Path: configured_directory = Path(settings.yolo_models_dir).expanduser() if configured_directory.exists() and configured_directory.is_dir(): return configured_directory.resolve() active_model_path = ModelAssetCatalogService._resolved_file_path(settings.yolo_model_path) if active_model_path and active_model_path.parent.exists() and active_model_path.parent.is_dir(): return active_model_path.parent.resolve() return configured_directory.resolve() @staticmethod def _asset_from_file(path: Path, *, active_model_path: Path | None) -> ModelAssetRead: resolved_path = path.resolve() return ModelAssetRead( model_asset_id=ModelAssetCatalogService._asset_id(path), filename=path.name, display_name=path.stem, model_path=str(resolved_path), suffix=path.suffix.lower(), framework=ModelAssetCatalogService.SUPPORTED_SUFFIXES[path.suffix.lower()], task_type="object_detection", size_bytes=path.stat().st_size, sha256=ModelAssetCatalogService._sha256(path), active=active_model_path == resolved_path, status="approved" if active_model_path == resolved_path else "available", limitation_message=( "Approved local runtime model asset. GeoIntel will not download or mutate model weights." if active_model_path == resolved_path else "Local development model asset. Configure it explicitly before production use." ), will_download_models=False, ) @staticmethod def _asset_id(path: Path) -> str: raw = f"{path.stem}-{path.suffix.lower().lstrip('.')}" normalized = re.sub(r"[^a-z0-9]+", "-", raw.lower()).strip("-") return normalized or "model-asset" @staticmethod def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() @staticmethod def _resolved_file_path(raw_path: str | None) -> Path | None: if not raw_path: return None path = Path(raw_path).expanduser() if not path.exists() or not path.is_file(): return None return path.resolve()