ylff/models/capture_models.py

"""
Pydantic models for Capture Bundle (SPECIFICATIONS.md Appendix C).

These are intentionally strict about structure, but permissive about forward
compatibility (extra fields allowed where reasonable) so we can evolve the schema.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field

from .spec_enums import OperatingRegime


class CaptureBundleSchemaVersion(str, Enum):
    V1_0 = "1.0"
    V2_0 = "2.0"


class DeviceType(str, Enum):
    IPHONE = "iphone"


class CaptureDeviceRef(BaseModel):
    """A device entry inside a capture bundle."""

    device_id: str = Field(..., description="Unique device identifier within bundle")
    device_type: DeviceType = Field(DeviceType.IPHONE, description="Device type")
    label: Optional[str] = Field(None, description="Human label (e.g. iphone_a/iphone_b) for rigs")

    # Paths are stored relative to bundle root to keep bundles relocatable.
    video_path: str = Field(..., description="Relative path to device video file")
    intrinsics_path: str = Field(..., description="Relative path to intrinsics json")
    timestamps_path: str = Field(..., description="Relative path to timestamps json")
    arkit_poses_path: Optional[str] = Field(
        None, description="Relative path to ARKit/VIO poses json (optional)"
    )
    lidar_depth_dir: Optional[str] = Field(
        None, description="Relative path to directory of depth frames (optional)"
    )

    model_config = {"extra": "allow"}


class CalibrationRef(BaseModel):
    rig_extrinsics_path: Optional[str] = Field(
        None, description="Relative path to rig extrinsics json"
    )
    sync_offsets_path: Optional[str] = Field(
        None, description="Relative path to per-device sync offsets json"
    )

    model_config = {"extra": "allow"}


class AnnotationRefs(BaseModel):
    quick_annotation_path: Optional[str] = Field(
        None, description="Relative path to quick annotation json"
    )
    detailed_annotation_path: Optional[str] = Field(
        None, description="Relative path to detailed annotation json"
    )

    model_config = {"extra": "allow"}


class TeacherOutputRefs(BaseModel):
    depth_dir: Optional[str] = Field(None, description="Relative path to teacher depth dir")
    uncertainty_dir: Optional[str] = Field(
        None, description="Relative path to teacher uncertainty dir"
    )
    reconstruction_path: Optional[str] = Field(
        None, description="Relative path to reconstruction artifact (ply/obj)"
    )

    model_config = {"extra": "allow"}


class CaptureManifest(BaseModel):
    """
    Canonical manifest for a capture bundle.

    NOTE: Appendix C in the spec lists a directory structure, but not the exact
    manifest schema. We define a stable minimal schema here and allow extras.
    """

    schema_version: CaptureBundleSchemaVersion = Field(CaptureBundleSchemaVersion.V1_0)
    capture_id: str = Field(..., description="Capture bundle ID")
    created_at: Optional[datetime] = Field(None, description="Capture creation timestamp")

    devices: List[CaptureDeviceRef] = Field(default_factory=list)
    calibration: Optional[CalibrationRef] = None
    annotations: Optional[AnnotationRefs] = None
    teacher_outputs: Optional[TeacherOutputRefs] = None

    # Optional metadata useful for stratification/regime selection.
    operating_regime: Optional[OperatingRegime] = None
    scene_type: Optional[str] = None
    difficulty_flags: List[str] = Field(default_factory=list)

    # Free-form metadata passthrough.
    metadata: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class AnnotationConfidence(str, Enum):
    CERTAIN = "certain"
    PROBABLE = "probable"
    UNSURE = "unsure"


class AnnotationObject(BaseModel):
    id: str
    type: str
    frame_time: float
    bbox: List[float] = Field(..., min_length=4, max_length=4)

    model_config = {"extra": "allow"}


class AnnotationSegment(BaseModel):
    id: str
    start_time: float
    end_time: float
    scene_type: str
    confidence: AnnotationConfidence = AnnotationConfidence.CERTAIN
    flags: List[str] = Field(default_factory=list)
    objects: List[AnnotationObject] = Field(default_factory=list)

    model_config = {"extra": "allow"}


class DetailedAnnotation(BaseModel):
    schema_version: Literal["1.0"] = "1.0"
    capture_id: str
    annotator_id: Optional[str] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None
    segments: List[AnnotationSegment] = Field(default_factory=list)
    scene_metadata: Dict[str, Any] = Field(default_factory=dict)
    quality_assessment: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}