geointel/backend/app/services/yolo_adapter.py

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)