Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |