Initial public release

backend/app/services/segmentation_adapter.py

@@ -0,0 +1,241 @@
+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