305 lines
9.9 KiB
Python
305 lines
9.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from app.core.config import Settings
|
|
from app.models import AnalysisRun, Dataset, Detection, Job, Project
|
|
from app.services.detection_georeferencing import pixel_bbox_to_epsg4326_polygon
|
|
from app.services.detection_service import DetectionService
|
|
from app.services.model_registry_service import ModelRegistryService
|
|
|
|
|
|
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 AvailableAdapter:
|
|
@staticmethod
|
|
def dependencies_available() -> bool:
|
|
return True
|
|
|
|
|
|
class MissingDependencyAdapter:
|
|
@staticmethod
|
|
def dependencies_available() -> bool:
|
|
return False
|
|
|
|
|
|
class MockYoloAdapter:
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
self.loaded_model_path: Path | None = None
|
|
|
|
@staticmethod
|
|
def dependencies_available() -> bool:
|
|
return True
|
|
|
|
def load_model(self, model_path: Path):
|
|
self.loaded_model_path = model_path
|
|
return object()
|
|
|
|
def predict_tile(self, model, tile_path: Path, confidence_threshold: float) -> list[dict]:
|
|
assert tile_path.name == "tile_0000.tif"
|
|
assert confidence_threshold == 0.5
|
|
return [
|
|
{
|
|
"class_name": "building",
|
|
"confidence": 0.91,
|
|
"bbox": [10.0, 20.0, 30.0, 40.0],
|
|
"properties": {"adapter": "mock"},
|
|
}
|
|
]
|
|
|
|
|
|
def _project_and_dataset(dataset_type: str = "raster"):
|
|
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=dataset_type,
|
|
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 _settings(tmp_path: Path, **overrides) -> Settings:
|
|
model_path = tmp_path / "model.pt"
|
|
values = {
|
|
"yolo_enabled": True,
|
|
"yolo_model_path": str(model_path),
|
|
"yolo_max_tiles": 4,
|
|
}
|
|
values.update(overrides)
|
|
return Settings(**values)
|
|
|
|
|
|
def _manifest(tmp_path: Path, tile_count: int = 1) -> Path:
|
|
tiles = []
|
|
for index in range(tile_count):
|
|
tile_path = tmp_path / f"tile_{index:04d}.tif"
|
|
tile_path.write_bytes(b"fixture")
|
|
tiles.append(
|
|
{
|
|
"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": index,
|
|
}
|
|
)
|
|
manifest_path = tmp_path / "manifest.json"
|
|
manifest_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"tile_set_id": "tiles-fixture",
|
|
"source_dataset_id": str(uuid4()),
|
|
"source_raster_id": str(uuid4()),
|
|
"tile_size": 100,
|
|
"overlap": 0,
|
|
"count": tile_count,
|
|
"tiles": tiles,
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
return manifest_path
|
|
|
|
|
|
def test_yolo_configured_model_reports_not_configured_when_disabled(tmp_path: Path) -> None:
|
|
settings = _settings(tmp_path, yolo_enabled=False)
|
|
|
|
models = {model.model_id: model for model in ModelRegistryService.list_model_capabilities(settings=settings)}
|
|
|
|
assert "yolo-configured" in models
|
|
assert models["yolo-configured"].configured is False
|
|
assert models["yolo-configured"].status == "not_configured"
|
|
|
|
|
|
def test_yolo_configured_model_reports_dependency_unavailable(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"local weights")
|
|
settings = _settings(tmp_path, yolo_model_path=str(model_path))
|
|
|
|
model = ModelRegistryService.get_model_capability(
|
|
"yolo-configured",
|
|
settings=settings,
|
|
yolo_adapter_class=MissingDependencyAdapter,
|
|
)
|
|
|
|
assert model is not None
|
|
assert model.configured is False
|
|
assert model.status == "dependency_unavailable"
|
|
|
|
|
|
def test_yolo_configured_model_reports_configured_with_local_model_and_dependencies(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"local weights")
|
|
settings = _settings(tmp_path, yolo_model_path=str(model_path))
|
|
|
|
model = ModelRegistryService.get_model_capability("yolo-configured", settings=settings, yolo_adapter_class=AvailableAdapter)
|
|
|
|
assert model is not None
|
|
assert model.configured is True
|
|
assert model.status == "configured"
|
|
assert model.version == settings.yolo_model_version
|
|
|
|
|
|
def test_yolo_run_requires_tile_manifest_path(tmp_path: Path) -> None:
|
|
db, project_id, dataset_id = _project_and_dataset()
|
|
settings = _settings(tmp_path)
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
DetectionService.run_detection(
|
|
db=db,
|
|
project_id=project_id,
|
|
dataset_id=dataset_id,
|
|
model_id="yolo-configured",
|
|
confidence_threshold=0.5,
|
|
settings=settings,
|
|
yolo_adapter_class=AvailableAdapter,
|
|
)
|
|
|
|
assert getattr(exc_info.value, "code", None) == "DETECTION_TILE_MANIFEST_REQUIRED"
|
|
|
|
|
|
def test_yolo_run_rejects_manifest_over_tile_limit(tmp_path: Path) -> None:
|
|
db, project_id, dataset_id = _project_and_dataset()
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"local weights")
|
|
settings = _settings(tmp_path, yolo_model_path=str(model_path), yolo_max_tiles=1)
|
|
manifest_path = _manifest(tmp_path, tile_count=2)
|
|
|
|
result = DetectionService.run_detection(
|
|
db=db,
|
|
project_id=project_id,
|
|
dataset_id=dataset_id,
|
|
model_id="yolo-configured",
|
|
confidence_threshold=0.5,
|
|
tile_manifest_path=str(manifest_path),
|
|
settings=settings,
|
|
yolo_adapter_class=MockYoloAdapter,
|
|
)
|
|
|
|
assert result.status == "failed"
|
|
assert result.error_code == "DETECTION_TILE_LIMIT_EXCEEDED"
|
|
|
|
|
|
def test_yolo_run_rejects_missing_tile_manifest_file(tmp_path: Path) -> None:
|
|
db, project_id, dataset_id = _project_and_dataset()
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"local weights")
|
|
settings = _settings(tmp_path, yolo_model_path=str(model_path))
|
|
|
|
result = DetectionService.run_detection(
|
|
db=db,
|
|
project_id=project_id,
|
|
dataset_id=dataset_id,
|
|
model_id="yolo-configured",
|
|
confidence_threshold=0.5,
|
|
tile_manifest_path=str(tmp_path / "missing-manifest.json"),
|
|
settings=settings,
|
|
yolo_adapter_class=MockYoloAdapter,
|
|
)
|
|
|
|
assert result.status == "failed"
|
|
assert result.error_code == "DETECTION_TILE_MANIFEST_NOT_FOUND"
|
|
|
|
|
|
def test_yolo_run_rejects_invalid_tile_manifest_json(tmp_path: Path) -> None:
|
|
db, project_id, dataset_id = _project_and_dataset()
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"local weights")
|
|
settings = _settings(tmp_path, yolo_model_path=str(model_path))
|
|
manifest_path = tmp_path / "manifest.json"
|
|
manifest_path.write_text("{not-json", encoding="utf-8")
|
|
|
|
result = DetectionService.run_detection(
|
|
db=db,
|
|
project_id=project_id,
|
|
dataset_id=dataset_id,
|
|
model_id="yolo-configured",
|
|
confidence_threshold=0.5,
|
|
tile_manifest_path=str(manifest_path),
|
|
settings=settings,
|
|
yolo_adapter_class=MockYoloAdapter,
|
|
)
|
|
|
|
assert result.status == "failed"
|
|
assert result.error_code == "DETECTION_TILE_MANIFEST_INVALID"
|
|
|
|
|
|
def test_pixel_bbox_to_epsg4326_polygon_from_gdal_transform() -> None:
|
|
polygon = pixel_bbox_to_epsg4326_polygon(
|
|
bbox=[10, 20, 30, 40],
|
|
tile={
|
|
"transform": [4.0, 0.01, 0.0, 52.0, 0.0, -0.01],
|
|
"bounds": [4.0, 51.0, 5.0, 52.0],
|
|
},
|
|
crs="EPSG:4326",
|
|
)
|
|
|
|
assert polygon.bounds == pytest.approx((4.1, 51.6, 4.3, 51.8))
|
|
|
|
|
|
def test_yolo_run_persists_mocked_georeferenced_detections(tmp_path: Path) -> None:
|
|
db, project_id, dataset_id = _project_and_dataset()
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"local weights")
|
|
settings = _settings(tmp_path, yolo_model_path=str(model_path), yolo_model_version="local-test")
|
|
manifest_path = _manifest(tmp_path, tile_count=1)
|
|
|
|
result = DetectionService.run_detection(
|
|
db=db,
|
|
project_id=project_id,
|
|
dataset_id=dataset_id,
|
|
model_id="yolo-configured",
|
|
confidence_threshold=0.5,
|
|
class_filter=["building"],
|
|
tile_manifest_path=str(manifest_path),
|
|
settings=settings,
|
|
yolo_adapter_class=MockYoloAdapter,
|
|
)
|
|
|
|
detections = [item for item in db.added if isinstance(item, Detection)]
|
|
runs = [item for item in db.added if isinstance(item, AnalysisRun)]
|
|
jobs = [item for item in db.added if isinstance(item, Job)]
|
|
|
|
assert result.status == "success"
|
|
assert result.detection_count == 1
|
|
assert detections[0].model_name == "yolo-configured"
|
|
assert detections[0].model_version == "local-test"
|
|
assert detections[0].class_name == "building"
|
|
assert detections[0].confidence == 0.91
|
|
assert detections[0].source_tile_path.endswith("tile_0000.tif")
|
|
assert detections[0].bbox_json == {"x_min": 10.0, "y_min": 20.0, "x_max": 30.0, "y_max": 40.0}
|
|
assert detections[0].properties_json == {"adapter": "mock", "tile_index": 0}
|
|
assert runs[0].status == "success"
|
|
assert jobs[0].status == "success"
|