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, ) 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