src/setup/models.py

"""
Pydantic models for the setup/template building phase.

These models are used during the initial template collection phase where
OCR-labeled samples are collected and averaged into digit templates.
Also includes region configuration models for setup.
"""

from typing import Any

import numpy as np
from pydantic import BaseModel, ConfigDict, Field


class DigitSample(BaseModel):
    """A single digit sample extracted from a play clock region."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    digit_value: int = Field(..., description="0-9 for ones digit, 0-4 for tens digit, -1 for blank")
    is_tens_digit: bool = Field(..., description="True if this is the tens place digit")
    position: str = Field(..., description="'left', 'center', or 'right' - where digit appears in region")
    image: np.ndarray[Any, Any] = Field(..., description="The digit image (grayscale, preprocessed)")
    source_clock_value: int = Field(..., description="The full clock value this was extracted from")
    timestamp: float = Field(..., description="Video timestamp where this was captured")
    confidence: float = Field(..., description="OCR confidence for this sample")


class DigitTemplate(BaseModel):
    """A template for matching a specific digit."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    digit_value: int = Field(..., description="0-9 for ones, 0-4 for tens, -1 for blank")
    is_tens_digit: bool = Field(..., description="True if this is a tens place template")
    position: str = Field(..., description="'left', 'center', or 'right' - where digit appears in region")
    template: np.ndarray[Any, Any] = Field(..., description="The template image (grayscale)")
    sample_count: int = Field(..., description="Number of samples used to build this template")
    avg_confidence: float = Field(..., description="Average OCR confidence of source samples")


class PlayClockRegionConfig(BaseModel):
    """Configuration for the play clock region relative to the scorebug bounding box."""

    x_offset: int = Field(..., description="X offset from scorebug left edge")
    y_offset: int = Field(..., description="Y offset from scorebug top edge")
    width: int = Field(..., description="Width of play clock region")
    height: int = Field(..., description="Height of play clock region")
    source_video: str = Field(..., description="Video used to identify region")
    scorebug_template: str = Field(..., description="Template used for scorebug detection")
    samples_used: int = Field(..., description="Number of frames used to verify region")