105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
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,
|
|
)
|
|
results = model.predict(
|
|
source=str(tile_path),
|
|
conf=float(confidence_threshold),
|
|
imgsz=int(self.settings.yolo_image_size),
|
|
device=self.settings.yolo_device,
|
|
verbose=False,
|
|
)
|
|
|
|
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)
|