""" Digit template builder for creating play clock digit templates from OCR samples. This module provides the DigitTemplateBuilder class which collects samples from the play clock region, extracts individual digits, and builds averaged templates for each unique digit value. Region extraction and preprocessing utilities are shared from utils to eliminate code duplication. """ import logging from typing import Any, Dict, List, Optional, Tuple import cv2 import numpy as np from utils import ( extract_center_region, extract_far_left_region, extract_left_region, extract_right_region, preprocess_playclock_region, ) from .coverage import ONES_DIGITS, categorize_template_keys from .models import DigitSample, DigitTemplate from .template_library import DigitTemplateLibrary logger = logging.getLogger(__name__) class DigitTemplateBuilder: """ Builds digit templates from OCR-labeled play clock samples. Collects samples from the play clock region, extracts individual digits, and builds averaged templates for each unique digit value. Uses color normalization so red and white digits produce the same template. """ # Play clock region dimensions (from config) DEFAULT_REGION_WIDTH = 50 DEFAULT_REGION_HEIGHT = 28 def __init__(self, region_width: int = DEFAULT_REGION_WIDTH, region_height: int = DEFAULT_REGION_HEIGHT): """ Initialize the template builder. Args: region_width: Width of play clock region in pixels region_height: Height of play clock region in pixels """ self.region_width = region_width self.region_height = region_height # Collected samples: {(is_tens, digit_value, position): [DigitSample, ...]} self.samples: Dict[Tuple[bool, int, str], List[DigitSample]] = {} # Track raw clock region images for potential reprocessing self.raw_regions: List[Tuple[float, int, np.ndarray[Any, Any]]] = [] # (timestamp, clock_value, region) logger.info("DigitTemplateBuilder initialized (region: %dx%d)", region_width, region_height) def preprocess_region(self, region: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]: """ Preprocess play clock region for template extraction. Delegates to shared utility function in utils.regions. Args: region: Play clock region (BGR format) Returns: Preprocessed binary image (white digits on black background) """ return preprocess_playclock_region(region, scale_factor=4) def extract_digits( self, preprocessed: np.ndarray[Any, Any], clock_value: int ) -> Tuple[Optional[np.ndarray[Any, Any]], Optional[np.ndarray[Any, Any]], Optional[np.ndarray[Any, Any]], Optional[np.ndarray[Any, Any]]]: """ Extract individual digit images from preprocessed play clock region. For double-digit values (10-40): extracts left (tens) and right (ones) For single-digit values (0-9): extracts far-left (blank) and center (ones) Args: preprocessed: Preprocessed play clock image (scaled 4x) clock_value: The known clock value (0-40) Returns: Tuple of (tens_digit_image, ones_right_image, ones_center_image, blank_image) - For double-digit: tens=left, ones_right=right, ones_center=None, blank=None - For single-digit: tens=None, ones_right=None, ones_center=center, blank=far_left """ if clock_value >= 10: # Double-digit: standard left/right split return extract_left_region(preprocessed), extract_right_region(preprocessed), None, None # Single-digit: far-left is blank (truly empty), ones is centered return None, None, extract_center_region(preprocessed), extract_far_left_region(preprocessed) def add_sample(self, region: np.ndarray[Any, Any], clock_value: int, timestamp: float, confidence: float = 1.0) -> None: """ Add a play clock sample for template building. Routes samples based on display layout: - Single-digit (0-9): Digit is CENTER-aligned, tens position is blank - Double-digit (10-40): Tens on LEFT, ones on RIGHT Args: region: Play clock region (BGR format, original size) clock_value: OCR-determined clock value (0-40) timestamp: Video timestamp confidence: OCR confidence score """ if clock_value < 0 or clock_value > 40: logger.warning("Invalid clock value %d, skipping sample", clock_value) return # Store raw region for potential reprocessing self.raw_regions.append((timestamp, clock_value, region.copy())) # Preprocess (handles red-to-white conversion automatically) preprocessed = self.preprocess_region(region) # Extract digits based on single vs double digit display tens_img, ones_right_img, ones_center_img, blank_img = self.extract_digits(preprocessed, clock_value) # Determine digit values ones_digit = clock_value % 10 tens_digit = clock_value // 10 if clock_value >= 10 else -1 # -1 = blank if clock_value >= 10: # Double-digit display (10-40): tens on left, ones on right assert tens_img is not None # Asserts: validated by extract_digits # Store tens sample (left position) tens_sample = DigitSample( digit_value=tens_digit, is_tens_digit=True, position="left", image=tens_img, source_clock_value=clock_value, timestamp=timestamp, confidence=confidence, ) tens_key = (True, tens_digit, "left") if tens_key not in self.samples: self.samples[tens_key] = [] self.samples[tens_key].append(tens_sample) # Store ones sample (right position) assert ones_right_img is not None # Asserts: validated by extract_digits ones_sample = DigitSample( digit_value=ones_digit, is_tens_digit=False, position="right", image=ones_right_img, source_clock_value=clock_value, timestamp=timestamp, confidence=confidence, ) ones_key = (False, ones_digit, "right") if ones_key not in self.samples: self.samples[ones_key] = [] self.samples[ones_key].append(ones_sample) logger.debug( "Added double-digit sample: clock=%d, tens=%d (left), ones=%d (right), t=%.1f", clock_value, tens_digit, ones_digit, timestamp, ) else: # Single-digit display (0-9): digit is centered, tens position is blank # Store blank sample (far-left position - should be truly empty) assert blank_img is not None # Asserts: validated by extract_digits blank_sample = DigitSample( digit_value=-1, # blank is_tens_digit=True, position="left", # Still use "left" as the position key for compatibility image=blank_img, # Now using far-left region that's truly empty source_clock_value=clock_value, timestamp=timestamp, confidence=confidence, ) blank_key = (True, -1, "left") if blank_key not in self.samples: self.samples[blank_key] = [] self.samples[blank_key].append(blank_sample) # Store ones sample (center position) assert ones_center_img is not None # Asserts: validated by extract_digits ones_sample = DigitSample( digit_value=ones_digit, is_tens_digit=False, position="center", image=ones_center_img, source_clock_value=clock_value, timestamp=timestamp, confidence=confidence, ) ones_key = (False, ones_digit, "center") if ones_key not in self.samples: self.samples[ones_key] = [] self.samples[ones_key].append(ones_sample) logger.debug( "Added single-digit sample: clock=%d, ones=%d (center), blank (far-left), t=%.1f", clock_value, ones_digit, timestamp, ) def get_sample_count(self) -> Dict[str, int]: """Get count of samples collected for each digit and position.""" counts = {} for (is_tens, digit, position), samples in self.samples.items(): type_str = "tens" if is_tens else "ones" digit_str = "blank" if digit == -1 else str(digit) key = f"{type_str}_{digit_str}_{position}" counts[key] = len(samples) return counts def build_templates(self, min_samples: int = 3) -> DigitTemplateLibrary: """ Build templates from collected samples. For each digit/position combination, averages multiple samples to create a robust template. Args: min_samples: Minimum samples required to build a template (default: 3) Returns: DigitTemplateLibrary with built templates """ library = DigitTemplateLibrary() for (is_tens, digit, position), samples in self.samples.items(): if len(samples) < min_samples: digit_display = "blank" if digit == -1 else str(digit) logger.warning( "Insufficient samples for %s digit %s (%s): %d < %d", "tens" if is_tens else "ones", digit_display, position, len(samples), min_samples, ) continue # Resize all samples to match dimensions of first sample target_shape = samples[0].image.shape # Average the samples (with resizing if needed) sum_image = np.zeros(target_shape, dtype=np.float32) valid_count = 0 total_confidence = 0.0 for sample in samples: img = sample.image if img.shape != target_shape: img = cv2.resize(img, (target_shape[1], target_shape[0])) sum_image += img.astype(np.float32) valid_count += 1 total_confidence += sample.confidence if valid_count > 0: avg_image = (sum_image / valid_count).astype(np.uint8) # Threshold the averaged image to clean it up _, template_img = cv2.threshold(avg_image, 127, 255, cv2.THRESH_BINARY) template = DigitTemplate( digit_value=digit, is_tens_digit=is_tens, position=position, template=template_img, sample_count=valid_count, avg_confidence=total_confidence / valid_count, ) library.add_template(template) digit_display = "blank" if digit == -1 else str(digit) logger.info( "Built template: %s digit %s (%s) from %d samples", "tens" if is_tens else "ones", digit_display, position, valid_count, ) # Log coverage status coverage = library.get_coverage_status() logger.info( "Template coverage: %d/%d (%.1f%%)", coverage["total_have"], coverage["total_needed"], 100 * coverage["total_have"] / coverage["total_needed"], ) return library def get_coverage_status(self) -> Dict[str, Any]: """Get current sample coverage status.""" # Get keys for samples that have at least one entry keys_with_samples = [key for key, samples in self.samples.items() if len(samples) >= 1] # Use shared utility to categorize ones_center_have, ones_right_have, tens_have, has_blank = categorize_template_keys(keys_with_samples) return { "ones_center": sorted(ones_center_have), "ones_right": sorted(ones_right_have), "tens": sorted(tens_have), "has_blank": has_blank, "ones_center_missing": sorted(ONES_DIGITS - ones_center_have), "ones_right_missing": sorted(ONES_DIGITS - ones_right_have), "tens_missing": sorted({1, 2, 3, 4} - tens_have), } def get_coverage_estimate(self) -> float: """ Get a simple coverage estimate as a float (0.0-1.0). Returns: Coverage estimate where 1.0 = all templates have samples """ status = self.get_coverage_status() # Count what we have (with at least 1 sample each) total_have = len(status["ones_center"]) + len(status["ones_right"]) + len(status["tens"]) if status["has_blank"]: total_have += 1 # Total needed: 10 ones_center + 10 ones_right + 4 tens + 1 blank = 25 total_needed = 25 return total_have / total_needed