""" 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]