Pre_Seg_Server/backend/models.py

"""SQLAlchemy ORM models."""

from sqlalchemy import (
    Column,
    Integer,
    String,
    Text,
    DateTime,
    ForeignKey,
    JSON,
    Float,
)
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func

from database import Base
from statuses import PROJECT_STATUS_PENDING


class Project(Base):
    """Project model representing a segmentation project."""

    __tablename__ = "projects"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(255), nullable=False)
    description = Column(Text, nullable=True)
    video_path = Column(String(512), nullable=True)
    thumbnail_url = Column(String(512), nullable=True)
    status = Column(String(50), default=PROJECT_STATUS_PENDING, nullable=False)
    source_type = Column(String(20), default="video", nullable=False)  # video | dicom
    original_fps = Column(Float, nullable=True)
    parse_fps = Column(Float, default=30.0, nullable=False)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )

    frames = relationship("Frame", back_populates="project", cascade="all, delete-orphan")
    annotations = relationship(
        "Annotation", back_populates="project", cascade="all, delete-orphan"
    )
    tasks = relationship(
        "ProcessingTask", back_populates="project", cascade="all, delete-orphan"
    )


class Frame(Base):
    """Frame model representing an extracted video frame."""

    __tablename__ = "frames"

    id = Column(Integer, primary_key=True, index=True)
    project_id = Column(Integer, ForeignKey("projects.id", ondelete="CASCADE"), nullable=False)
    frame_index = Column(Integer, nullable=False)
    image_url = Column(String(512), nullable=False)
    width = Column(Integer, nullable=True)
    height = Column(Integer, nullable=True)
    timestamp_ms = Column(Float, nullable=True)
    source_frame_number = Column(Integer, nullable=True)
    created_at = Column(DateTime(timezone=True), server_default=func.now())

    project = relationship("Project", back_populates="frames")
    annotations = relationship(
        "Annotation", back_populates="frame", cascade="all, delete-orphan"
    )


class Template(Base):
    """Template (Ontology) model for segmentation classes."""

    __tablename__ = "templates"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(255), nullable=False)
    description = Column(Text, nullable=True)
    color = Column(String(50), nullable=False)
    z_index = Column(Integer, default=0, nullable=False)
    mapping_rules = Column(JSON, nullable=True)
    created_at = Column(DateTime(timezone=True), server_default=func.now())

    annotations = relationship(
        "Annotation", back_populates="template", cascade="all, delete-orphan"
    )


class Annotation(Base):
    """Annotation model for segmentation masks and prompts."""

    __tablename__ = "annotations"

    id = Column(Integer, primary_key=True, index=True)
    project_id = Column(
        Integer, ForeignKey("projects.id", ondelete="CASCADE"), nullable=False
    )
    frame_id = Column(
        Integer, ForeignKey("frames.id", ondelete="CASCADE"), nullable=True
    )
    template_id = Column(
        Integer, ForeignKey("templates.id", ondelete="SET NULL"), nullable=True
    )
    mask_data = Column(JSON, nullable=True)
    points = Column(JSON, nullable=True)
    bbox = Column(JSON, nullable=True)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )

    project = relationship("Project", back_populates="annotations")
    frame = relationship("Frame", back_populates="annotations")
    template = relationship("Template", back_populates="annotations")
    masks = relationship("Mask", back_populates="annotation", cascade="all, delete-orphan")


class Mask(Base):
    """Mask model for exported/derived mask files."""

    __tablename__ = "masks"

    id = Column(Integer, primary_key=True, index=True)
    annotation_id = Column(
        Integer, ForeignKey("annotations.id", ondelete="CASCADE"), nullable=False
    )
    mask_url = Column(String(512), nullable=False)
    format = Column(String(50), default="png", nullable=False)  # png / rle / json
    created_at = Column(DateTime(timezone=True), server_default=func.now())

    annotation = relationship("Annotation", back_populates="masks")


class ProcessingTask(Base):
    """Background task state persisted for dashboard and polling."""

    __tablename__ = "processing_tasks"

    id = Column(Integer, primary_key=True, index=True)
    task_type = Column(String(80), nullable=False)
    status = Column(String(40), default="queued", nullable=False)
    progress = Column(Integer, default=0, nullable=False)
    message = Column(Text, nullable=True)
    project_id = Column(
        Integer, ForeignKey("projects.id", ondelete="CASCADE"), nullable=True
    )
    celery_task_id = Column(String(255), nullable=True)
    payload = Column(JSON, nullable=True)
    result = Column(JSON, nullable=True)
    error = Column(Text, nullable=True)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    started_at = Column(DateTime(timezone=True), nullable=True)
    finished_at = Column(DateTime(timezone=True), nullable=True)
    updated_at = Column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )

    project = relationship("Project", back_populates="tasks")