117 lines
5.1 KiB
Python
117 lines
5.1 KiB
Python
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,
|
|
runtime_available=True,
|
|
runtime_status="active" if active_model_path == resolved_path else "available",
|
|
governed_validation_status="not_verified_by_catalog",
|
|
promotion_status="not_verified_by_catalog",
|
|
status="runtime_active" if active_model_path == resolved_path else "runtime_available",
|
|
limitation_message=(
|
|
"Active local runtime model asset. Runtime selection is not evidence of governed validation or promotion."
|
|
if active_model_path == resolved_path
|
|
else "Local runtime model asset. Governed validation and promotion are not established by this catalog."
|
|
),
|
|
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()
|