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, ) self.validate_runtime() 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 validate_runtime(self) -> None: if not self.settings.yolo_require_cuda: return try: import torch except Exception as exc: raise AppError( code="DETECTION_ACCELERATOR_UNAVAILABLE", message="NVIDIA CUDA is required for configured YOLO inference, but PyTorch is not importable.", status_code=503, ) from exc if not torch.cuda.is_available(): raise AppError( code="DETECTION_ACCELERATOR_UNAVAILABLE", message="NVIDIA CUDA is required for configured YOLO inference, but no CUDA device is available.", details={"configured_device": self.settings.yolo_device}, status_code=503, ) if not str(self.settings.yolo_device).lower().startswith(("cuda", "0", "1", "2", "3")): raise AppError( code="DETECTION_ACCELERATOR_MISCONFIGURED", message="NVIDIA CUDA is required, but YOLO_DEVICE does not select a CUDA device.", details={"configured_device": self.settings.yolo_device}, status_code=503, ) 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)