151 lines
5.3 KiB
Python
151 lines
5.3 KiB
Python
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,
|
|
)
|
|
|
|
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 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,
|
|
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)
|