geointel/backend/app/services/segmentation_adapter.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

from app.core.config import Settings
from app.core.errors import AppError
from app.services.yolo_adapter import _prediction_source, _to_list


@dataclass(frozen=True)
class SegmentationAdapterResult:
    class_name: str
    confidence: float | None
    geometry: dict[str, Any]
    bbox_json: dict[str, Any] | None = None
    mask_path: str | None = None
    source_tile_path: str | None = None
    tile_index: int | None = None
    properties_json: dict[str, Any] | None = None
    provenance_json: dict[str, Any] | None = None
    area_m2: float | None = None


class SegmentationAdapter(Protocol):
    def segment(self, *args: Any, **kwargs: Any) -> list[SegmentationAdapterResult]:
        """Future segmentation adapters must local-import model dependencies inside execution paths."""


class _UltralyticsSegmentationAdapterBase:
    """Shared local-inference plumbing for ultralytics-backed segmentation models.

    Model weights are never downloaded automatically; a missing local file or
    missing dependency fails closed with an explicit error.
    """

    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 _require_model_file(self, model_path: Path) -> None:
        if not model_path.exists() or not model_path.is_file():
            raise AppError(
                code="SEGMENTATION_MODEL_UNAVAILABLE",
                message="Configured segmentation model file does not exist",
                details={"model_path": str(model_path)},
                status_code=503,
            )
        if not self.dependencies_available():
            raise AppError(
                code="SEGMENTATION_DEPENDENCY_UNAVAILABLE",
                message="Segmentation dependencies are not installed. Install backend optional extras with geointel-backend[ai].",
                status_code=503,
            )
        self.validate_runtime()

    def validate_runtime(self) -> None:
        """Fail closed when the deployment contract requires NVIDIA CUDA.

        Detection and segmentation share ``YOLO_DEVICE`` and
        ``YOLO_REQUIRE_CUDA``. Without this check segmentation could advertise
        a GPU job while Ultralytics silently used CPU or failed only after the
        model had already been loaded.
        """

        if not self.settings.yolo_require_cuda:
            return
        try:
            import torch
        except Exception as exc:
            raise AppError(
                code="SEGMENTATION_ACCELERATOR_UNAVAILABLE",
                message="NVIDIA CUDA is required for configured segmentation, but PyTorch is not importable.",
                status_code=503,
            ) from exc
        if not torch.cuda.is_available():
            raise AppError(
                code="SEGMENTATION_ACCELERATOR_UNAVAILABLE",
                message="NVIDIA CUDA is required for configured segmentation, 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="SEGMENTATION_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(self, model, tile_path: Path, confidence_threshold: float) -> list[Any]:
        if not tile_path.exists() or not tile_path.is_file():
            raise AppError(
                code="SEGMENTATION_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:
                return 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="SEGMENTATION_INFERENCE_FAILED",
                message="Configured segmentation inference failed for a raster tile",
                details={"tile_path": str(tile_path), "error": str(exc)},
                status_code=503,
            ) from exc

    def _extract_masks(self, results: list[Any], default_class_name: str | None = None) -> list[dict[str, Any]]:
        segmentations: list[dict[str, Any]] = []
        max_masks = int(self.settings.segmentation_max_masks_per_tile)
        for result in results:
            names = getattr(result, "names", {}) or {}
            masks = getattr(result, "masks", None)
            if masks is None:
                continue
            polygons = getattr(masks, "xy", None) or []
            boxes = getattr(result, "boxes", None)
            confidence_values = _to_list(getattr(boxes, "conf", [])) if boxes is not None else []
            class_values = _to_list(getattr(boxes, "cls", [])) if boxes is not None else []
            bbox_values = _to_list(getattr(boxes, "xyxy", [])) if boxes is not None else []
            for index, polygon in enumerate(polygons):
                if len(segmentations) >= max_masks:
                    return segmentations
                points = _to_list(polygon)
                if not isinstance(points, list) or len(points) < 3:
                    continue
                class_id = int(class_values[index]) if index < len(class_values) else -1
                if default_class_name is not None:
                    class_name = default_class_name
                else:
                    class_name = str(names.get(class_id, class_id))
                confidence = float(confidence_values[index]) if index < len(confidence_values) else None
                bbox = [float(value) for value in bbox_values[index]] if index < len(bbox_values) else None
                segmentations.append(
                    {
                        "class_name": class_name,
                        "confidence": confidence,
                        "points": [[float(point[0]), float(point[1])] for point in points],
                        "bbox": bbox,
                        "properties": {"class_id": class_id},
                    }
                )
        return segmentations


class YoloSegmentationAdapter(_UltralyticsSegmentationAdapterBase):
    def load_model(self, model_path: Path):
        self._require_model_file(model_path)
        try:
            from ultralytics import YOLO
        except ImportError as exc:
            raise AppError(
                code="SEGMENTATION_DEPENDENCY_UNAVAILABLE",
                message="YOLO segmentation 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="SEGMENTATION_MODEL_LOAD_FAILED",
                message="Configured YOLO segmentation 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]]:
        results = self._predict(model, tile_path, confidence_threshold)
        return self._extract_masks(results)


class SamSegmentationAdapter(_UltralyticsSegmentationAdapterBase):
    """Class-agnostic SAM segmentation through the ultralytics SAM interface."""

    def load_model(self, model_path: Path):
        self._require_model_file(model_path)
        try:
            from ultralytics import SAM
        except ImportError as exc:
            raise AppError(
                code="SEGMENTATION_DEPENDENCY_UNAVAILABLE",
                message="SAM segmentation requires the ultralytics SAM interface. Install backend optional extras with geointel-backend[ai].",
                status_code=503,
            ) from exc
        try:
            return SAM(str(model_path))
        except Exception as exc:
            raise AppError(
                code="SEGMENTATION_MODEL_LOAD_FAILED",
                message="Configured SAM 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]]:
        results = self._predict(model, tile_path, confidence_threshold)
        return self._extract_masks(results, default_class_name="segment")


class FixtureSegmentationAdapter:
    def segment(self, raw_segmentations: Any) -> list[SegmentationAdapterResult]:
        if not isinstance(raw_segmentations, list):
            return []
        results: list[SegmentationAdapterResult] = []
        for raw in raw_segmentations:
            if not isinstance(raw, dict):
                continue
            results.append(
                SegmentationAdapterResult(
                    class_name=str(raw.get("class_name") or ""),
                    confidence=float(raw["confidence"]) if raw.get("confidence") is not None else None,
                    geometry=raw.get("geometry"),
                    bbox_json=raw.get("bbox_json"),
                    mask_path=raw.get("mask_path"),
                    source_tile_path=raw.get("source_tile_path"),
                    tile_index=raw.get("tile_index"),
                    properties_json=raw.get("properties_json"),
                    provenance_json=raw.get("provenance_json"),
                    area_m2=raw.get("area_m2"),
                )
            )
        return results