Add local model asset catalog
This commit is contained in:
@@ -80,6 +80,7 @@ def test_env_example_uses_runtime_env_names_read_by_backend_and_frontend() -> No
|
||||
|
||||
assert "GEOINTEL_INSTALL_AI=false" in env_example
|
||||
assert "YOLO_ENABLED=false" in env_example
|
||||
assert "YOLO_MODELS_DIR=/app/models" in env_example
|
||||
assert "YOLO_MODEL_PATH=" in env_example
|
||||
assert "YOLO_CONFIG_DIR=./storage/ultralytics" in env_example
|
||||
assert "YOLO_MAX_TILES=100" in env_example
|
||||
@@ -231,6 +232,8 @@ def test_unraid_deploy_passes_ai_build_arg_and_yolo_runtime_env() -> None:
|
||||
|
||||
assert 'YOLO_ENABLED="${YOLO_ENABLED:-false}"' in run_script
|
||||
assert '-e YOLO_ENABLED="$YOLO_ENABLED"' in run_script
|
||||
assert 'YOLO_MODELS_DIR="${YOLO_MODELS_DIR:-/app/models}"' in run_script
|
||||
assert '-e YOLO_MODELS_DIR="$YOLO_MODELS_DIR"' in run_script
|
||||
assert '-e YOLO_MODEL_PATH="$YOLO_MODEL_PATH"' in run_script
|
||||
assert '-e YOLO_MAX_TILES="$YOLO_MAX_TILES"' in run_script
|
||||
assert "-v \"${GEOINTEL_MODELS_PATH}:/app/models\"" in run_script
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.core.config import Settings
|
||||
from app.core.errors import AppError
|
||||
from app.main import app
|
||||
from app.models import AnalysisRun, Dataset, Detection, Job, Project
|
||||
from app.services.detection_service import DetectionService
|
||||
from app.services.model_asset_catalog_service import ModelAssetCatalogService
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, objects=None) -> None:
|
||||
self.objects = objects or {}
|
||||
self.added = []
|
||||
self.commits = 0
|
||||
self.refreshes = []
|
||||
|
||||
def get(self, model, item_id):
|
||||
return self.objects.get((model, item_id))
|
||||
|
||||
def add(self, item) -> None:
|
||||
self.added.append(item)
|
||||
if getattr(item, "id", None) is not None:
|
||||
self.objects[(item.__class__, item.id)] = item
|
||||
|
||||
def commit(self) -> None:
|
||||
self.commits += 1
|
||||
|
||||
def refresh(self, item) -> None:
|
||||
self.refreshes.append(item)
|
||||
|
||||
|
||||
class MockYoloAdapter:
|
||||
def __init__(self, settings: Settings) -> None:
|
||||
self.settings = settings
|
||||
|
||||
@staticmethod
|
||||
def dependencies_available() -> bool:
|
||||
return True
|
||||
|
||||
def load_model(self, model_path: Path):
|
||||
return {"model_path": str(model_path)}
|
||||
|
||||
def predict_tile(self, model, tile_path: Path, confidence_threshold: float) -> list[dict]:
|
||||
assert model["model_path"].endswith("building-detector.pt")
|
||||
return [
|
||||
{
|
||||
"class_name": "building",
|
||||
"confidence": 0.9,
|
||||
"bbox": [10.0, 20.0, 30.0, 40.0],
|
||||
"properties": {"adapter": "mock"},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _project_and_raster_dataset():
|
||||
project_id = uuid4()
|
||||
dataset_id = uuid4()
|
||||
project = Project(id=project_id, name="Geel")
|
||||
dataset = Dataset(
|
||||
id=dataset_id,
|
||||
project_id=project_id,
|
||||
name="source.tif",
|
||||
dataset_type="raster",
|
||||
source="user_upload",
|
||||
storage_path="storage/uploads/source.tif",
|
||||
)
|
||||
db = FakeSession(objects={(Project, project_id): project, (Dataset, dataset_id): dataset})
|
||||
return db, project_id, dataset_id
|
||||
|
||||
|
||||
def _manifest(tmp_path: Path) -> Path:
|
||||
tile_path = tmp_path / "tile_0000.tif"
|
||||
tile_path.write_bytes(b"tile")
|
||||
manifest_path = tmp_path / "manifest.json"
|
||||
manifest_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"tile_set_id": "tiles-fixture",
|
||||
"count": 1,
|
||||
"tiles": [
|
||||
{
|
||||
"path": str(tile_path),
|
||||
"pixel_window": [0, 0, 100, 100],
|
||||
"bounds": [4.0, 51.0, 5.0, 52.0],
|
||||
"transform": [4.0, 0.01, 0.0, 52.0, 0.0, -0.01],
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest_path
|
||||
|
||||
|
||||
def test_model_asset_catalog_lists_supported_local_model_files(tmp_path: Path) -> None:
|
||||
model_file = tmp_path / "building-detector.pt"
|
||||
model_file.write_bytes(b"local model")
|
||||
ignored_file = tmp_path / "notes.txt"
|
||||
ignored_file.write_text("ignore me", encoding="utf-8")
|
||||
settings = Settings(yolo_models_dir=str(tmp_path), yolo_model_path=str(model_file), yolo_enabled=True)
|
||||
|
||||
response = ModelAssetCatalogService.list_assets(settings=settings)
|
||||
|
||||
assert response.total == 1
|
||||
asset = response.items[0]
|
||||
assert asset.model_asset_id == "building-detector-pt"
|
||||
assert asset.filename == "building-detector.pt"
|
||||
assert asset.display_name == "building-detector"
|
||||
assert asset.model_path == str(model_file)
|
||||
assert asset.size_bytes == len(b"local model")
|
||||
assert len(asset.sha256) == 64
|
||||
assert asset.active is True
|
||||
assert asset.status == "available"
|
||||
assert asset.will_download_models is False
|
||||
|
||||
|
||||
def test_model_asset_catalog_resolves_known_asset(tmp_path: Path) -> None:
|
||||
model_file = tmp_path / "building-detector.pt"
|
||||
model_file.write_bytes(b"local model")
|
||||
settings = Settings(yolo_models_dir=str(tmp_path), yolo_enabled=True)
|
||||
|
||||
asset = ModelAssetCatalogService.resolve_asset("building-detector-pt", settings=settings)
|
||||
|
||||
assert asset.filename == "building-detector.pt"
|
||||
assert asset.model_path == str(model_file)
|
||||
|
||||
|
||||
def test_model_asset_catalog_rejects_unknown_asset(tmp_path: Path) -> None:
|
||||
settings = Settings(yolo_models_dir=str(tmp_path), yolo_enabled=True)
|
||||
|
||||
with pytest.raises(AppError) as exc_info:
|
||||
ModelAssetCatalogService.resolve_asset("missing-model", settings=settings)
|
||||
|
||||
assert exc_info.value.code == "DETECTION_MODEL_ASSET_NOT_FOUND"
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
def test_model_assets_api_returns_canonical_envelope(monkeypatch, tmp_path: Path) -> None:
|
||||
model_file = tmp_path / "building-detector.pt"
|
||||
model_file.write_bytes(b"local model")
|
||||
monkeypatch.setenv("YOLO_MODELS_DIR", str(tmp_path))
|
||||
monkeypatch.setenv("YOLO_MODEL_PATH", str(model_file))
|
||||
|
||||
response = TestClient(app).get("/api/v1/detection/model-assets")
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert set(payload) == {"data"}
|
||||
assert payload["data"]["total"] == 1
|
||||
assert payload["data"]["items"][0]["model_asset_id"] == "building-detector-pt"
|
||||
assert payload["data"]["items"][0]["active"] is True
|
||||
assert payload["data"]["items"][0]["will_download_models"] is False
|
||||
|
||||
|
||||
def test_detection_run_persists_selected_model_asset_parameters(tmp_path: Path) -> None:
|
||||
model_file = tmp_path / "building-detector.pt"
|
||||
model_file.write_bytes(b"local model")
|
||||
db, project_id, dataset_id = _project_and_raster_dataset()
|
||||
settings = Settings(
|
||||
yolo_enabled=True,
|
||||
yolo_model_path=str(tmp_path / "default.pt"),
|
||||
yolo_models_dir=str(tmp_path),
|
||||
yolo_max_tiles=4,
|
||||
)
|
||||
|
||||
result = DetectionService.run_detection(
|
||||
db=db,
|
||||
project_id=project_id,
|
||||
dataset_id=dataset_id,
|
||||
model_id="yolo-configured",
|
||||
model_asset_id="building-detector-pt",
|
||||
confidence_threshold=0.5,
|
||||
tile_manifest_path=str(_manifest(tmp_path)),
|
||||
settings=settings,
|
||||
yolo_adapter_class=MockYoloAdapter,
|
||||
)
|
||||
|
||||
jobs = [item for item in db.added if isinstance(item, Job)]
|
||||
runs = [item for item in db.added if isinstance(item, AnalysisRun)]
|
||||
detections = [item for item in db.added if isinstance(item, Detection)]
|
||||
|
||||
assert result.status == "success"
|
||||
assert result.detection_count == 1
|
||||
assert jobs[0].parameters_json["model_asset_id"] == "building-detector-pt"
|
||||
assert jobs[0].parameters_json["model_asset_path"] == str(model_file)
|
||||
assert len(jobs[0].parameters_json["model_asset_sha256"]) == 64
|
||||
assert runs[0].parameters_json["model_asset_id"] == "building-detector-pt"
|
||||
assert detections[0].model_name == "yolo-configured"
|
||||
@@ -23,3 +23,29 @@ def test_detection_lab_surfaces_yolo_runtime_preflight() -> None:
|
||||
assert "/api/v1/detection/yolo/preflight" in api
|
||||
assert "interface YoloPreflightResponse" in types
|
||||
assert "yoloPreflight={yoloPreflight}" in app
|
||||
|
||||
|
||||
def test_detection_lab_surfaces_local_model_asset_selection() -> None:
|
||||
lab = (ROOT / "frontend" / "src" / "components" / "detection" / "DetectionLab.tsx").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
hook = (ROOT / "frontend" / "src" / "hooks" / "useDetectionWorkflow.ts").read_text(encoding="utf-8")
|
||||
api = (ROOT / "frontend" / "src" / "services" / "api" / "detection.ts").read_text(encoding="utf-8")
|
||||
types = (ROOT / "frontend" / "src" / "types.ts").read_text(encoding="utf-8")
|
||||
app = (ROOT / "frontend" / "src" / "App.tsx").read_text(encoding="utf-8")
|
||||
provider_panel = (ROOT / "frontend" / "src" / "components" / "providers" / "ProviderPanel.tsx").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
assert "interface ModelAssetRead" in types
|
||||
assert "model_asset_id?: string | null" in types
|
||||
assert "listModelAssets" in api
|
||||
assert "/api/v1/detection/model-assets" in api
|
||||
assert "modelAssets" in hook
|
||||
assert "selectedModelAssetId" in hook
|
||||
assert "model_asset_id: selectedModelAssetId || null" in hook
|
||||
assert "Local model assets" in lab
|
||||
assert "onSelectModelAsset" in lab
|
||||
assert "modelAssets={modelAssets}" in app
|
||||
assert "Official reference sources" in provider_panel
|
||||
assert "not AI model choices" in provider_panel
|
||||
|
||||
@@ -66,6 +66,7 @@ def test_configure_yolo_model_dry_run_selects_single_model(tmp_path: Path) -> No
|
||||
assert payload["selected_container_model_path"] == "/app/models/nested/detector.pt"
|
||||
assert payload["env_updates"]["GEOINTEL_INSTALL_AI"] == "true"
|
||||
assert payload["env_updates"]["YOLO_ENABLED"] == "true"
|
||||
assert payload["env_updates"]["YOLO_MODELS_DIR"] == "/app/models"
|
||||
assert payload["env_updates"]["YOLO_MODEL_PATH"] == "/app/models/nested/detector.pt"
|
||||
assert payload["will_download_models"] is False
|
||||
assert not (tmp_path / ".env").exists()
|
||||
@@ -94,4 +95,5 @@ def test_configure_yolo_model_apply_updates_existing_env_file(tmp_path: Path) ->
|
||||
assert "GEOINTEL_FRONTEND_PORT=1202" in contents
|
||||
assert "GEOINTEL_INSTALL_AI=true" in contents
|
||||
assert "YOLO_ENABLED=true" in contents
|
||||
assert "YOLO_MODELS_DIR=/app/models" in contents
|
||||
assert "YOLO_MODEL_PATH=/app/models/detector.engine" in contents
|
||||
|
||||
Reference in New Issue
Block a user