src/ui/models.py

"""
Pydantic models and data structures for the UI module.

This module contains all the data models used for region selection,
including bounding boxes, view configurations, and selection state.
"""

from typing import Any, List, Optional, Tuple

import numpy as np
from pydantic import BaseModel, ConfigDict


class BBox(BaseModel):
    """Bounding box with x, y, width, height coordinates."""

    x: int
    y: int
    width: int
    height: int

    @property
    def x2(self) -> int:
        """Get the right edge x coordinate."""
        return self.x + self.width

    @property
    def y2(self) -> int:
        """Get the bottom edge y coordinate."""
        return self.y + self.height

    def to_tuple(self) -> Tuple[int, int, int, int]:
        """Convert to (x, y, width, height) tuple."""
        return (self.x, self.y, self.width, self.height)

    def scaled(self, factor: int) -> "BBox":
        """Return a new BBox scaled by the given factor."""
        return BBox(x=self.x * factor, y=self.y * factor, width=self.width * factor, height=self.height * factor)

    def unscaled(self, factor: int) -> "BBox":
        """Return a new BBox divided by the given scale factor."""
        return BBox(x=self.x // factor, y=self.y // factor, width=self.width // factor, height=self.height // factor)

    def offset(self, dx: int, dy: int) -> "BBox":
        """Return a new BBox offset by dx, dy."""
        return BBox(x=self.x + dx, y=self.y + dy, width=self.width, height=self.height)

    @classmethod
    def from_points(cls, p1: Tuple[int, int], p2: Tuple[int, int]) -> "BBox":
        """Create BBox from two corner points."""
        x = min(p1[0], p2[0])
        y = min(p1[1], p2[1])
        w = abs(p2[0] - p1[0])
        h = abs(p2[1] - p1[1])
        return cls(x=x, y=y, width=w, height=h)

    @classmethod
    def from_tuple(cls, t: Tuple[int, int, int, int]) -> "BBox":
        """Create BBox from (x, y, width, height) tuple."""
        return cls(x=t[0], y=t[1], width=t[2], height=t[3])


class SelectionViewConfig(BaseModel):
    """Configuration for display/view settings during selection."""

    scale_factor: int = 1
    padding: int = 0
    window_width: int = 1280
    window_height: int = 720


class SelectionState(BaseModel):
    """Mutable state for a region selection session."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    frame_idx: int = 0
    frames: List[Tuple[float, np.ndarray[Any, Any]]] = []
    should_quit: bool = False
    should_confirm: bool = False
    video_path: Optional[str] = None

    @property
    def current_frame(self) -> Tuple[float, np.ndarray[Any, Any]]:
        """Get the current (timestamp, frame) tuple."""
        return self.frames[self.frame_idx]

    @property
    def timestamp(self) -> float:
        """Get the current frame timestamp."""
        return self.frames[self.frame_idx][0]

    @property
    def frame(self) -> np.ndarray[Any, Any]:
        """Get the current frame."""
        return self.frames[self.frame_idx][1]