Jens
faeb58ef6d
GeoIntel release gates / Compile, test, contracts and builds (push) Successful in 1m49s
GeoIntel release gates / Python and npm vulnerability policy (push) Successful in 21s
GeoIntel release gates / Production AI image, SBOM and container scan (push) Successful in 5m39s
GeoIntel release gates / Deploy exact gated revision to Unraid (push) Failing after 58m43s
105 lines
3.1 KiB
Python
105 lines
3.1 KiB
Python
"""Tiles must reach the GPU in batches.
|
|
|
|
``YOLO_BATCH_SIZE`` existed in the settings but nothing read it: every tile was
|
|
a separate ``model.predict`` call plus a separate temporary PNG. On an RTX-class
|
|
card that leaves most of the throughput unused for a run of a hundred tiles.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from app.core.errors import AppError
|
|
from app.services.yolo_adapter import YoloDetectionAdapter
|
|
|
|
|
|
class RecordingModel:
|
|
def __init__(self) -> None:
|
|
self.batches: list[list[str]] = []
|
|
|
|
def predict(self, *, source, conf, imgsz, device, verbose, max_det):
|
|
self.batches.append(list(source) if isinstance(source, list) else [source])
|
|
return []
|
|
|
|
|
|
def _settings(tmp_path: Path, **overrides):
|
|
from app.core.config import Settings
|
|
|
|
values = {
|
|
"yolo_model_path": str(tmp_path / "model.pt"),
|
|
"yolo_device": "cpu",
|
|
"yolo_image_size": 64,
|
|
"yolo_max_detections": 1000,
|
|
"yolo_require_cuda": False,
|
|
"yolo_batch_size": 4,
|
|
}
|
|
values.update(overrides)
|
|
return Settings(**values)
|
|
|
|
|
|
def _tiles(tmp_path: Path, count: int) -> list[Path]:
|
|
Image = pytest.importorskip("PIL.Image")
|
|
paths = []
|
|
for index in range(count):
|
|
path = tmp_path / f"tile_{index:04d}.png"
|
|
Image.new("RGB", (16, 16), (index, 20, 30)).save(path)
|
|
paths.append(path)
|
|
return paths
|
|
|
|
|
|
def test_tiles_are_predicted_in_configured_batches(tmp_path: Path) -> None:
|
|
tiles = _tiles(tmp_path, 9)
|
|
model = RecordingModel()
|
|
|
|
YoloDetectionAdapter(_settings(tmp_path, yolo_batch_size=4)).predict_tiles(model, tiles, 0.25)
|
|
|
|
assert [len(batch) for batch in model.batches] == [4, 4, 1]
|
|
|
|
|
|
def test_batch_size_one_still_works(tmp_path: Path) -> None:
|
|
tiles = _tiles(tmp_path, 3)
|
|
model = RecordingModel()
|
|
|
|
YoloDetectionAdapter(_settings(tmp_path, yolo_batch_size=1)).predict_tiles(model, tiles, 0.25)
|
|
|
|
assert [len(batch) for batch in model.batches] == [1, 1, 1]
|
|
|
|
|
|
def test_results_are_returned_per_tile_in_order(tmp_path: Path) -> None:
|
|
tiles = _tiles(tmp_path, 3)
|
|
|
|
class PerTileModel:
|
|
def predict(self, *, source, conf, imgsz, device, verbose, max_det):
|
|
sources = list(source) if isinstance(source, list) else [source]
|
|
return [_FakeResult(index) for index, _ in enumerate(sources)]
|
|
|
|
results = YoloDetectionAdapter(_settings(tmp_path)).predict_tiles(PerTileModel(), tiles, 0.25)
|
|
|
|
assert len(results) == 3
|
|
assert [len(detections) for detections in results] == [1, 1, 1]
|
|
|
|
|
|
def test_a_missing_tile_is_reported_before_the_batch_runs(tmp_path: Path) -> None:
|
|
tiles = _tiles(tmp_path, 2) + [tmp_path / "absent.png"]
|
|
|
|
with pytest.raises(AppError) as exc_info:
|
|
YoloDetectionAdapter(_settings(tmp_path)).predict_tiles(RecordingModel(), tiles, 0.25)
|
|
|
|
assert exc_info.value.code == "DETECTION_TILE_NOT_FOUND"
|
|
|
|
|
|
class _FakeBoxes:
|
|
xyxy = [[0.0, 0.0, 4.0, 4.0]]
|
|
conf = [0.9]
|
|
cls = [0]
|
|
|
|
|
|
class _FakeResult:
|
|
names = {0: "building"}
|
|
boxes = _FakeBoxes()
|
|
|
|
def __init__(self, _index: int) -> None:
|
|
pass
|