Pre_Seg_Server/backend/schemas.py

"""Pydantic schemas for request/response validation."""

from datetime import datetime
from typing import Optional, Any
from pydantic import BaseModel, ConfigDict


# ---------------------------------------------------------------------------
# Project schemas
# ---------------------------------------------------------------------------
class ProjectBase(BaseModel):
    name: str
    description: Optional[str] = None
    video_path: Optional[str] = None
    thumbnail_url: Optional[str] = None
    status: Optional[str] = "pending"
    source_type: Optional[str] = "video"
    original_fps: Optional[float] = None
    parse_fps: Optional[float] = 30.0


class ProjectCreate(ProjectBase):
    pass


class ProjectUpdate(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None
    video_path: Optional[str] = None
    thumbnail_url: Optional[str] = None
    status: Optional[str] = None
    source_type: Optional[str] = None
    original_fps: Optional[float] = None
    parse_fps: Optional[float] = None


class ProjectOut(ProjectBase):
    model_config = ConfigDict(from_attributes=True)

    id: int
    created_at: datetime
    updated_at: datetime
    frame_count: int = 0


# ---------------------------------------------------------------------------
# Frame schemas
# ---------------------------------------------------------------------------
class FrameBase(BaseModel):
    frame_index: int
    image_url: str
    width: Optional[int] = None
    height: Optional[int] = None
    timestamp_ms: Optional[float] = None
    source_frame_number: Optional[int] = None


class FrameCreate(FrameBase):
    project_id: int


class FrameOut(FrameBase):
    model_config = ConfigDict(from_attributes=True)

    id: int
    project_id: int
    created_at: datetime


# ---------------------------------------------------------------------------
# Template schemas
# ---------------------------------------------------------------------------
class TemplateBase(BaseModel):
    name: str
    description: Optional[str] = None
    color: str
    z_index: int = 0
    mapping_rules: Optional[dict[str, Any]] = None
    classes: Optional[list[dict[str, Any]]] = None
    rules: Optional[list[dict[str, Any]]] = None


class TemplateCreate(TemplateBase):
    pass


class TemplateUpdate(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None
    color: Optional[str] = None
    z_index: Optional[int] = None
    mapping_rules: Optional[dict[str, Any]] = None
    classes: Optional[list[dict[str, Any]]] = None
    rules: Optional[list[dict[str, Any]]] = None


class TemplateOut(TemplateBase):
    model_config = ConfigDict(from_attributes=True)

    id: int
    created_at: datetime


# ---------------------------------------------------------------------------
# Annotation schemas
# ---------------------------------------------------------------------------
class AnnotationBase(BaseModel):
    project_id: int
    frame_id: Optional[int] = None
    template_id: Optional[int] = None
    mask_data: Optional[dict[str, Any]] = None
    points: Optional[list[list[float]]] = None
    bbox: Optional[list[float]] = None


class AnnotationCreate(AnnotationBase):
    pass


class AnnotationUpdate(BaseModel):
    mask_data: Optional[dict[str, Any]] = None
    points: Optional[list[list[float]]] = None
    bbox: Optional[list[float]] = None
    template_id: Optional[int] = None


class AnnotationOut(AnnotationBase):
    model_config = ConfigDict(from_attributes=True)

    id: int
    created_at: datetime
    updated_at: datetime


# ---------------------------------------------------------------------------
# Mask schemas
# ---------------------------------------------------------------------------
class MaskBase(BaseModel):
    annotation_id: int
    mask_url: str
    format: str = "png"


class MaskCreate(MaskBase):
    pass


class MaskOut(MaskBase):
    model_config = ConfigDict(from_attributes=True)

    id: int
    created_at: datetime


# ---------------------------------------------------------------------------
# Processing task schemas
# ---------------------------------------------------------------------------
class ProcessingTaskOut(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: int
    task_type: str
    status: str
    progress: int
    message: Optional[str] = None
    project_id: Optional[int] = None
    celery_task_id: Optional[str] = None
    payload: Optional[dict[str, Any]] = None
    result: Optional[dict[str, Any]] = None
    error: Optional[str] = None
    created_at: datetime
    started_at: Optional[datetime] = None
    finished_at: Optional[datetime] = None
    updated_at: datetime


# ---------------------------------------------------------------------------
# AI schemas
# ---------------------------------------------------------------------------
class PredictRequest(BaseModel):
    image_id: int
    prompt_type: str  # point / box / semantic
    prompt_data: Any
    model: Optional[str] = None
    options: Optional[dict[str, Any]] = None


class PredictResponse(BaseModel):
    polygons: list[list[list[float]]]
    scores: Optional[list[float]] = None


class MaskAnalysisRequest(BaseModel):
    frame_id: Optional[int] = None
    mask_data: dict[str, Any]
    points: Optional[list[list[float]]] = None
    bbox: Optional[list[float]] = None
    extract_skeleton: bool = False


class MaskAnalysisResponse(BaseModel):
    confidence: Optional[float] = None
    confidence_source: str
    topology_anchor_count: int
    topology_anchors: list[list[float]]
    area: float
    bbox: Optional[list[float]] = None
    source: Optional[str] = None
    message: str


class PropagationSeed(BaseModel):
    polygons: Optional[list[list[list[float]]]] = None
    bbox: Optional[list[float]] = None
    points: Optional[list[list[float]]] = None
    labels: Optional[list[int]] = None
    label: Optional[str] = None
    color: Optional[str] = None
    class_metadata: Optional[dict[str, Any]] = None
    template_id: Optional[int] = None


class PropagateRequest(BaseModel):
    project_id: int
    frame_id: int
    model: Optional[str] = "sam2.1_hiera_tiny"
    seed: PropagationSeed
    direction: str = "forward"
    max_frames: int = 30
    include_source: bool = False
    save_annotations: bool = True


class PropagateResponse(BaseModel):
    model: str
    direction: str
    source_frame_id: int
    processed_frame_count: int
    created_annotation_count: int
    annotations: list[AnnotationOut]


class AiModelStatus(BaseModel):
    id: str
    label: str
    available: bool
    loaded: bool = False
    device: str
    supports: list[str]
    message: str
    package_available: bool = False
    checkpoint_exists: bool = False
    checkpoint_path: Optional[str] = None
    python_ok: bool = True
    torch_ok: bool = True
    cuda_required: bool = False
    external_available: bool = False
    external_python: Optional[str] = None


class GpuStatus(BaseModel):
    available: bool
    device: str
    name: Optional[str] = None
    torch_available: bool
    torch_version: Optional[str] = None
    cuda_version: Optional[str] = None


class AiRuntimeStatus(BaseModel):
    selected_model: str
    gpu: GpuStatus
    models: list[AiModelStatus]


# ---------------------------------------------------------------------------
# Export schemas
# ---------------------------------------------------------------------------
class ExportStatus(BaseModel):
    url: str
    format: str