LabelingSystem/backend/schemas/export.py

"""
Pydantic schemas for Export API.
Defines request and response models for data export operations.
"""
from datetime import datetime
from typing import Optional, List, Any
from pydantic import BaseModel, Field
from enum import Enum


class ExportFormat(str, Enum):
    """Supported export formats."""
    JSON = "json"
    CSV = "csv"
    COCO = "coco"
    YOLO = "yolo"
    PASCAL_VOC = "pascal_voc"
    SHAREGPT = "sharegpt"
    ALPACA = "alpaca"


class ExportStatus(str, Enum):
    """Export job status."""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class StatusFilter(str, Enum):
    """Task status filter for export."""
    ALL = "all"
    COMPLETED = "completed"
    PENDING = "pending"
    IN_PROGRESS = "in_progress"


class ExportRequest(BaseModel):
    """Schema for export request."""
    
    format: ExportFormat = Field(..., description="Export format: json, csv, coco, yolo")
    status_filter: StatusFilter = Field(
        default=StatusFilter.ALL,
        description="Filter tasks by status: all, completed, pending, in_progress"
    )
    include_metadata: bool = Field(
        default=True,
        description="Whether to include task metadata in export"
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "format": "json",
                "status_filter": "completed",
                "include_metadata": True
            }
        }


class ExportJobResponse(BaseModel):
    """Schema for export job response."""
    
    id: str = Field(..., description="Export job unique identifier")
    project_id: str = Field(..., description="Project ID being exported")
    format: str = Field(..., description="Export format")
    status: str = Field(..., description="Export job status")
    status_filter: str = Field(..., description="Task status filter used")
    include_metadata: bool = Field(..., description="Whether metadata is included")
    file_path: Optional[str] = Field(None, description="Path to exported file")
    download_url: Optional[str] = Field(None, description="URL to download the file")
    error_message: Optional[str] = Field(None, description="Error message if failed")
    created_at: datetime = Field(..., description="Job creation timestamp")
    completed_at: Optional[datetime] = Field(None, description="Job completion timestamp")
    total_tasks: int = Field(default=0, description="Total number of tasks to export")
    exported_tasks: int = Field(default=0, description="Number of tasks exported")
    
    class Config:
        json_schema_extra = {
            "example": {
                "id": "export_abc123",
                "project_id": "proj_123abc",
                "format": "json",
                "status": "completed",
                "status_filter": "completed",
                "include_metadata": True,
                "file_path": "exports/export_abc123.json",
                "download_url": "/api/exports/export_abc123/download",
                "error_message": None,
                "created_at": "2024-01-12T12:00:00",
                "completed_at": "2024-01-12T12:01:00",
                "total_tasks": 100,
                "exported_tasks": 100
            }
        }


class ExportProgressResponse(BaseModel):
    """Schema for export progress response."""
    
    id: str = Field(..., description="Export job unique identifier")
    status: str = Field(..., description="Export job status")
    progress: float = Field(..., description="Export progress (0.0 to 1.0)")
    total_tasks: int = Field(..., description="Total number of tasks")
    exported_tasks: int = Field(..., description="Number of tasks exported")
    error_message: Optional[str] = Field(None, description="Error message if failed")


class AnnotationExportItem(BaseModel):
    """Schema for a single annotation in export."""
    
    id: str = Field(..., description="Annotation ID")
    task_id: str = Field(..., description="Task ID")
    user_id: str = Field(..., description="User ID who created the annotation")
    result: Any = Field(..., description="Annotation result data")
    created_at: str = Field(..., description="Creation timestamp")
    updated_at: str = Field(..., description="Last update timestamp")


class TaskExportItem(BaseModel):
    """Schema for a single task in export."""
    
    id: str = Field(..., description="Task ID")
    name: str = Field(..., description="Task name")
    data: Any = Field(..., description="Task data")
    status: str = Field(..., description="Task status")
    assigned_to: Optional[str] = Field(None, description="Assigned user ID")
    created_at: str = Field(..., description="Creation timestamp")
    annotations: List[AnnotationExportItem] = Field(
        default_factory=list,
        description="List of annotations for this task"
    )


class ProjectExportData(BaseModel):
    """Schema for complete project export data."""
    
    project_id: str = Field(..., description="Project ID")
    project_name: str = Field(..., description="Project name")
    project_description: Optional[str] = Field(None, description="Project description")
    config: str = Field(..., description="Project label config")
    export_format: str = Field(..., description="Export format used")
    export_time: str = Field(..., description="Export timestamp")
    total_tasks: int = Field(..., description="Total number of tasks exported")
    total_annotations: int = Field(..., description="Total number of annotations")
    tasks: List[TaskExportItem] = Field(..., description="List of exported tasks")


# COCO format schemas
class COCOImage(BaseModel):
    """COCO format image entry."""
    id: int
    file_name: str
    width: int = 0
    height: int = 0


class COCOCategory(BaseModel):
    """COCO format category entry."""
    id: int
    name: str
    supercategory: str = ""


class COCOAnnotation(BaseModel):
    """COCO format annotation entry."""
    id: int
    image_id: int
    category_id: int
    bbox: List[float] = Field(default_factory=list)
    area: float = 0.0
    segmentation: List[Any] = Field(default_factory=list)
    iscrowd: int = 0


class COCOExportData(BaseModel):
    """COCO format export data."""
    info: dict = Field(default_factory=dict)
    licenses: List[dict] = Field(default_factory=list)
    images: List[COCOImage] = Field(default_factory=list)
    annotations: List[COCOAnnotation] = Field(default_factory=list)
    categories: List[COCOCategory] = Field(default_factory=list)