""" Digit template library for storing and managing play clock digit templates. This module provides the DigitTemplateLibrary class for saving, loading, and managing digit templates used for play clock reading. """ import json import logging from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import cv2 from .coverage import ONES_DIGITS, calculate_coverage_status, categorize_template_keys from .models import DigitTemplate logger = logging.getLogger(__name__) class DigitTemplateLibrary: """ Stores and manages digit templates for play clock reading. Uses color normalization to handle both red and white digits with a single template set. Now supports position-aware templates to handle both single-digit (centered) and double-digit (left/right split) layouts: - Ones digits (center): 0-9 from single-digit displays (10 templates) - Ones digits (right): 0-9 from double-digit displays (10 templates) - Tens digits (left): 1, 2, 3, 4 from double-digit displays (4 templates) - Blank (left): Empty tens position from single-digit displays (1 template) Total: 25 templates needed for full coverage """ # Template coverage requirements ONES_DIGITS = list(range(10)) # 0-9 TENS_DIGITS = [-1, 1, 2, 3, 4] # -1 = blank, 1-4 for 10-40 POSITIONS = ["left", "center", "right"] def __init__(self) -> None: """Initialize empty template library.""" # Templates: {(is_tens, digit_value, position): DigitTemplate} self.templates: Dict[Tuple[bool, int, str], DigitTemplate] = {} logger.info("DigitTemplateLibrary initialized (empty)") def add_template(self, template: DigitTemplate) -> None: """ Add a template to the library. Args: template: DigitTemplate to add """ key = (template.is_tens_digit, template.digit_value, template.position) self.templates[key] = template digit_display = "blank" if template.digit_value == -1 else str(template.digit_value) logger.debug( "Added template: tens=%s, digit=%s, position=%s", template.is_tens_digit, digit_display, template.position, ) def get_template(self, is_tens: bool, digit_value: int, position: str) -> Optional[DigitTemplate]: """ Get a template from the library. Args: is_tens: Whether this is a tens digit digit_value: The digit value (-1 for blank, 0-9 for digits) position: Template position ("left", "center", or "right") Returns: DigitTemplate if found, None otherwise """ key = (is_tens, digit_value, position) return self.templates.get(key) def get_all_templates(self, is_tens: bool, position: Optional[str] = None) -> List[DigitTemplate]: """ Get all templates for a specific digit position. Args: is_tens: Whether to get tens digit templates position: Optional position filter ("left", "center", or "right") Returns: List of matching DigitTemplate objects """ templates = [] for (tens, _, pos), template in self.templates.items(): if tens == is_tens: if position is None or pos == position: templates.append(template) return templates def get_coverage_status(self) -> Dict[str, Any]: """ Get the current template coverage status. Returns: Dictionary with coverage information """ # Use shared utility to categorize template keys ones_center_have, ones_right_have, tens_have, has_blank = categorize_template_keys(self.templates.keys()) # Get base coverage status from shared utility status = calculate_coverage_status(ones_center_have, ones_right_have, tens_have, has_blank, total_items=len(self.templates)) # Convert -1 to "blank" for display def format_tens(digits: set[int]) -> list[str | int]: return sorted(["blank" if d == -1 else d for d in digits], key=lambda x: (isinstance(x, str), x)) # Add legacy fields for backward compatibility status["ones_have"] = sorted(ones_center_have | ones_right_have) status["ones_missing"] = sorted((ONES_DIGITS - ones_center_have) & (ONES_DIGITS - ones_right_have)) status["tens_have_formatted"] = format_tens(tens_have | ({-1} if has_blank else set())) status["tens_missing_formatted"] = format_tens(set(status["tens_missing"]) | (set() if has_blank else {-1})) return status def is_complete(self) -> bool: """Check if all required templates are present.""" return bool(self.get_coverage_status()["is_complete"]) def save(self, output_path: str) -> None: """ Save templates to disk. Args: output_path: Path to save directory """ output_dir = Path(output_path) output_dir.mkdir(parents=True, exist_ok=True) templates_list: list[dict[str, object]] = [] for (is_tens, digit, position), template in self.templates.items(): # Use "blank" instead of -1 for the empty tens digit in filenames digit_str = "blank" if digit == -1 else str(digit) position_suffix = f"_{position}" if position != "left" or not is_tens else "" filename = f"{'tens' if is_tens else 'ones'}_{digit_str}{position_suffix}.png" cv2.imwrite(str(output_dir / filename), template.template) templates_list.append( { "filename": filename, "digit_value": digit_str, # Use "blank" for display "is_tens_digit": template.is_tens_digit, "position": position, "sample_count": template.sample_count, "avg_confidence": template.avg_confidence, } ) metadata = {"templates": templates_list, "version": 2} # Version 2 includes position with open(output_dir / "templates_metadata.json", "w", encoding="utf-8") as f: json.dump(metadata, f, indent=2) logger.info("Saved %d templates to %s", len(self.templates), output_path) def load(self, input_path: str) -> bool: """ Load templates from disk. Args: input_path: Path to templates directory Returns: True if loaded successfully, False otherwise """ input_dir = Path(input_path) metadata_path = input_dir / "templates_metadata.json" if not metadata_path.exists(): logger.warning("No templates metadata found at %s", metadata_path) return False with open(metadata_path, "r", encoding="utf-8") as f: metadata = json.load(f) version = metadata.get("version", 1) for entry in metadata.get("templates", []): img_path = input_dir / entry["filename"] if img_path.exists(): template_img = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE) if template_img is not None: # Convert "blank" back to -1 for internal use digit_value = entry["digit_value"] if digit_value == "blank": digit_value = -1 elif isinstance(digit_value, str): digit_value = int(digit_value) # Handle position (v2) or infer from old format (v1) is_tens = entry["is_tens_digit"] if version >= 2: position = entry.get("position", "left" if is_tens else "right") else: # V1 format: tens → left, ones → right (old behavior) position = "left" if is_tens else "right" template = DigitTemplate( digit_value=digit_value, is_tens_digit=is_tens, position=position, template=template_img, sample_count=entry.get("sample_count", 1), avg_confidence=entry.get("avg_confidence", 1.0), ) self.add_template(template) logger.info("Loaded %d templates from %s (v%d format)", len(self.templates), input_path, version) return True