Harden YOLO tile inference inputs

backend/tests/test_sprint8b_yolo_foundation.py

@@ -7,10 +7,12 @@ from uuid import uuid4
 import pytest
 
 from app.core.config import Settings
+from app.core.errors import AppError
 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
+from app.services.yolo_adapter import YoloDetectionAdapter
 
 ROOT = Path(__file__).resolve().parents[2]
 
@@ -75,6 +77,33 @@ class MockYoloAdapter:
         ]
 
 
+class RecordingPredictModel:
+    def __init__(self) -> None:
+        self.seen_sources: list[dict] = []
+
+    def predict(self, *, source, conf, imgsz, device, verbose):
+        from PIL import Image
+
+        with Image.open(source) as image:
+            self.seen_sources.append(
+                {
+                    "path": str(source),
+                    "mode": image.mode,
+                    "bands": len(image.getbands()),
+                    "conf": conf,
+                    "imgsz": imgsz,
+                    "device": device,
+                    "verbose": verbose,
+                }
+            )
+        return []
+
+
+class ExplodingPredictModel:
+    def predict(self, *, source, conf, imgsz, device, verbose):
+        raise RuntimeError("expected input[1, 1, 480, 640] to have 3 channels, but got 1 channels instead")
+
+
 def _project_and_dataset(dataset_type: str = "raster"):
     project_id = uuid4()
     dataset_id = uuid4()
@@ -312,3 +341,35 @@ def test_yolo_run_persists_mocked_georeferenced_detections(tmp_path: Path) -> No
     assert detections[0].properties_json == {"adapter": "mock", "tile_index": 0}
     assert runs[0].status == "success"
     assert jobs[0].status == "success"
+
+
+def test_yolo_adapter_converts_single_band_tiles_to_rgb_before_prediction(tmp_path: Path) -> None:
+    Image = pytest.importorskip("PIL.Image")
+    tile_path = tmp_path / "single_band_tile.tif"
+    Image.new("L", (16, 16), 128).save(tile_path)
+    model = RecordingPredictModel()
+    settings = _settings(tmp_path, yolo_image_size=64, yolo_device="cpu")
+
+    detections = YoloDetectionAdapter(settings).predict_tile(model, tile_path, confidence_threshold=0.25)
+
+    assert detections == []
+    assert model.seen_sources[0]["mode"] == "RGB"
+    assert model.seen_sources[0]["bands"] == 3
+    assert model.seen_sources[0]["path"] != str(tile_path)
+    assert model.seen_sources[0]["conf"] == 0.25
+    assert model.seen_sources[0]["imgsz"] == 64
+    assert model.seen_sources[0]["device"] == "cpu"
+    assert model.seen_sources[0]["verbose"] is False
+
+
+def test_yolo_adapter_wraps_prediction_runtime_errors(tmp_path: Path) -> None:
+    tile_path = tmp_path / "tile.tif"
+    tile_path.write_bytes(b"not an image but present")
+    settings = _settings(tmp_path)
+
+    with pytest.raises(AppError) as exc_info:
+        YoloDetectionAdapter(settings).predict_tile(ExplodingPredictModel(), tile_path, confidence_threshold=0.25)
+
+    assert exc_info.value.code == "DETECTION_INFERENCE_FAILED"
+    assert "Configured YOLO inference failed for a raster tile" in exc_info.value.message
+    assert exc_info.value.details["tile_path"] == str(tile_path)