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 class YoloDetectionAdapter: def __init__(self, settings: Settings) -> None: self.settings = settings @staticmethod def dependencies_available() -> bool: try: import torch # noqa: F401 import ultralytics # noqa: F401 except Exception: return False return True def load_model(self, model_path: Path): if not model_path.exists() or not model_path.is_file(): raise AppError( code="DETECTION_MODEL_UNAVAILABLE", message="Configured YOLO model file does not exist", details={"model_path": str(model_path)}, status_code=503, ) if not self.dependencies_available(): raise AppError( code="DETECTION_DEPENDENCY_UNAVAILABLE", message="YOLO dependencies are not installed. Install backend optional extras with geointel-backend[ai].", status_code=503, ) try: from ultralytics import YOLO except ImportError as exc: raise AppError( code="DETECTION_DEPENDENCY_UNAVAILABLE", message="YOLO dependencies are not importable. Install backend optional extras with geointel-backend[ai].", status_code=503, ) from exc try: return YOLO(str(model_path)) except Exception as exc: raise AppError( code="DETECTION_MODEL_LOAD_FAILED", message="Configured YOLO model could not be loaded", details={"model_path": str(model_path)}, status_code=503, ) from exc def predict_tile(self, model, tile_path: Path, confidence_threshold: float) -> list[dict[str, Any]]: if not tile_path.exists() or not tile_path.is_file(): raise AppError( code="DETECTION_TILE_NOT_FOUND", message="Tile referenced by manifest does not exist", details={"tile_path": str(tile_path)}, status_code=422, ) 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, max_det=int(self.settings.yolo_max_detections), 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: names = getattr(result, "names", {}) or {} boxes = getattr(result, "boxes", None) if boxes is None: continue xyxy_values = _to_list(getattr(boxes, "xyxy", [])) confidence_values = _to_list(getattr(boxes, "conf", [])) class_values = _to_list(getattr(boxes, "cls", [])) for index, bbox in enumerate(xyxy_values): class_id = int(class_values[index]) if index < len(class_values) else -1 detections.append( { "class_name": str(names.get(class_id, class_id)), "confidence": float(confidence_values[index]) if index < len(confidence_values) else 0.0, "bbox": [float(value) for value in bbox], "properties": {"class_id": class_id}, } ) return detections def _to_list(value: Any) -> list[Any]: if hasattr(value, "detach"): value = value.detach() if hasattr(value, "cpu"): value = value.cpu() if hasattr(value, "numpy"): value = value.numpy() 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)