Initial public release

backend/tests/test_yolo_batched_inference.py

@@ -0,0 +1,104 @@
+"""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