geointel/backend/app/schemas/detection.py

from __future__ import annotations

from datetime import datetime
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field


class DetectionModelCapability(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    model_id: str
    display_name: str
    framework: str
    task_type: str
    supported_classes: list[str]
    configured: bool
    status: str
    limitation_message: str
    version: str | None = None
    training_scope: str | None = None
    validation_scope: str | None = None
    validated_regions: list[str] = Field(default_factory=list)
    nationally_validated: bool = False
    operator_review_required: bool = True


class DetectionModelsResponse(BaseModel):
    models: list[DetectionModelCapability]


class ModelAssetRead(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    model_asset_id: str
    filename: str
    display_name: str
    model_path: str
    suffix: str
    framework: str
    task_type: str
    size_bytes: int
    sha256: str
    active: bool
    status: str
    limitation_message: str
    will_download_models: bool = False


class ModelAssetListResponse(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    items: list[ModelAssetRead]
    total: int
    model_directory: str


class DetectionRunRequest(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    project_id: UUID
    dataset_id: UUID
    model_id: str
    model_asset_id: str | None = None
    confidence_threshold: float = Field(default=0.5, ge=0.0, le=1.0)
    class_filter: list[str] | None = None
    tile_manifest_path: str | None = None
    parameters_json: dict = Field(default_factory=dict)


class DetectionQaRequest(BaseModel):
    reference_dataset_id: UUID
    iou_threshold: float = Field(default=0.5, ge=0.0, le=1.0)
    class_name: str | None = None
    min_confidence: float | None = Field(default=None, ge=0.0, le=1.0)
    # Confidence cuts to report alongside the run's own operating point. They
    # are read off the one matching pass, so a sweep costs no extra inference.
    calibration_thresholds: list[float] = Field(default_factory=list, max_length=32)


class DetectionComparisonRequest(BaseModel):
    """Place several runs side by side against one reference."""

    analysis_run_ids: list[UUID] = Field(min_length=2, max_length=12)
    reference_dataset_id: UUID
    iou_threshold: float = Field(default=0.5, ge=0.0, le=1.0)


class DetectionComparisonResponse(BaseModel):
    reference_dataset_id: UUID
    iou_threshold: float
    # Whether these runs answer the same question at all, and why not if they
    # do not. Numbers from incomparable runs are reported but never ranked as
    # if they were alternatives.
    comparability: dict
    ranking_metric: str
    rows: list[dict]


class DetectionRunResponse(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    analysis_run_id: UUID
    job_id: UUID
    project_id: UUID
    dataset_id: UUID
    model_id: str
    status: str
    detection_count: int
    error_code: str | None = None
    message: str


class DetectionRunRead(BaseModel):
    model_config = ConfigDict(from_attributes=True, protected_namespaces=())

    id: UUID
    project_id: UUID
    dataset_id: UUID | None = None
    job_id: UUID | None = None
    analysis_type: str
    status: str
    model_name: str | None = None
    model_version: str | None = None
    parameters_json: dict
    result_json: dict | None = None
    error_message: str | None = None
    created_at: datetime | None = None
    started_at: datetime | None = None
    finished_at: datetime | None = None


class DetectionRunListResponse(BaseModel):
    items: list[DetectionRunRead]
    # ``total`` counts every run; ``items`` is the most recent page of them.
    total: int
    limit: int | None = None
    offset: int = 0
    truncated: bool = False


class DetectionRead(BaseModel):
    model_config = ConfigDict(from_attributes=True, protected_namespaces=())

    id: UUID
    project_id: UUID
    dataset_id: UUID | None = None
    analysis_run_id: UUID | None = None
    job_id: UUID | None = None
    model_name: str
    model_version: str | None = None
    class_name: str
    confidence: float
    bbox_json: dict | None = None
    source_tile_path: str | None = None
    properties_json: dict | None = None
    created_at: datetime | None = None


class DetectionListResponse(BaseModel):
    items: list[DetectionRead]
    # ``total`` is the complete population; ``items`` is one page of it.
    total: int
    limit: int | None = None
    offset: int = 0
    truncated: bool = False


class YoloPreflightChecks(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    enabled: bool
    dependencies_available: bool | None = None
    accelerator_ready: bool | None = None
    model_path_set: bool | None = None
    model_file_exists: bool | None = None
    model_provenance_manifest_path: str | None = None
    model_provenance_valid: bool | None = None
    model_load_requested: bool
    model_load_ok: bool | None = None
    manifest_path_set: bool | None = None
    manifest_valid: bool | None = None
    tile_paths_exist: bool | None = None
    tile_limit_ok: bool | None = None


class YoloRuntimeDetails(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    dependencies_assumed: bool
    model_directory: str | None = None
    yolo_config_dir: str | None = None
    torch_version: str | None = None
    ultralytics_version: str | None = None
    cuda_available: bool | None = None
    configured_device: str
    cuda_required: bool


class YoloPreflightResponse(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    model_id: str
    model_asset_id: str | None = None
    model_path: str | None = None
    tile_manifest_path: str | None = None
    status: str
    message: str
    checks: YoloPreflightChecks
    tile_count: int
    max_tiles: int
    will_download_models: bool
    will_run_inference: bool
    runtime: YoloRuntimeDetails
    error_code: str | None = None
    details: dict | None = None