Harden YOLO tile inference inputs
GeoIntel CI / docs-smoke (push) Has been cancelled
GeoIntel CI / contract-smoke (push) Has been cancelled

This commit is contained in:
Codex
2026-07-06 20:25:09 +02:00
parent d7f786729a
commit add4768a52
7 changed files with 137 additions and 12 deletions
+53 -7
View File
@@ -1,7 +1,10 @@
from __future__ import annotations
from contextlib import contextmanager
from pathlib import Path
import tempfile
from typing import Any
from collections.abc import Iterator
from app.core.config import Settings
from app.core.errors import AppError
@@ -62,13 +65,24 @@ class YoloDetectionAdapter:
details={"tile_path": str(tile_path)},
status_code=422,
)
results = model.predict(
source=str(tile_path),
conf=float(confidence_threshold),
imgsz=int(self.settings.yolo_image_size),
device=self.settings.yolo_device,
verbose=False,
)
try:
with _prediction_source(tile_path) as prediction_source:
results = model.predict(
source=prediction_source,
conf=float(confidence_threshold),
imgsz=int(self.settings.yolo_image_size),
device=self.settings.yolo_device,
verbose=False,
)
except AppError:
raise
except Exception as exc:
raise AppError(
code="DETECTION_INFERENCE_FAILED",
message="Configured YOLO inference failed for a raster tile",
details={"tile_path": str(tile_path), "error": str(exc)},
status_code=503,
) from exc
detections: list[dict[str, Any]] = []
for result in results:
@@ -102,3 +116,35 @@ def _to_list(value: Any) -> list[Any]:
if hasattr(value, "tolist"):
return value.tolist()
return list(value)
@contextmanager
def _prediction_source(tile_path: Path) -> Iterator[str]:
temp_path: Path | None = None
try:
try:
from PIL import Image
except Exception:
yield str(tile_path)
return
try:
with Image.open(tile_path) as image:
if image.mode == "RGB" and len(image.getbands()) == 3:
yield str(tile_path)
return
rgb_image = image.convert("RGB")
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as handle:
temp_path = Path(handle.name)
rgb_image.save(temp_path)
yield str(temp_path)
return
except Exception:
if temp_path is not None:
raise
yield str(tile_path)
return
finally:
if temp_path is not None:
temp_path.unlink(missing_ok=True)