Spaces:
Sleeping
Sleeping
| """ | |
| Play clock region extraction and OCR preprocessing for template building. | |
| This module provides: | |
| - PlayClockRegionExtractor: Extracts and preprocesses play clock regions | |
| - OCR preprocessing for initial digit labeling during template building | |
| The region extraction logic determines WHERE to look in the frame, | |
| while the OCR preprocessing prepares images for EasyOCR labeling. | |
| Color detection utilities are shared from utils.color to eliminate code duplication. | |
| """ | |
| import json | |
| import logging | |
| from pathlib import Path | |
| from typing import Any, Optional, Tuple | |
| import cv2 | |
| import numpy as np | |
| from utils import detect_red_digits | |
| from .models import PlayClockRegionConfig | |
| logger = logging.getLogger(__name__) | |
| # ============================================================================= | |
| # Play Clock Region Extractor | |
| # ============================================================================= | |
| class PlayClockRegionExtractor: | |
| """ | |
| Extracts and preprocesses play clock regions from video frames. | |
| The extractor locates the play clock sub-region within the scorebug | |
| and preprocesses it for OCR during template building. This class | |
| handles the geometry of WHERE to look in the frame. | |
| Note: For reading actual clock values, use ReadPlayClock from readers.playclock. | |
| """ | |
| def __init__(self, region_config_path: Optional[str] = None, region_config: Optional[PlayClockRegionConfig] = None): | |
| """ | |
| Initialize the play clock region extractor. | |
| Args: | |
| region_config_path: Path to JSON config file with play clock region coordinates | |
| region_config: Direct PlayClockRegionConfig object (alternative to file path) | |
| """ | |
| self.config: Optional[PlayClockRegionConfig] = None | |
| if region_config: | |
| self.config = region_config | |
| logger.info("PlayClockRegionExtractor initialized with direct config") | |
| elif region_config_path: | |
| self.load_config(region_config_path) | |
| else: | |
| logger.warning("PlayClockRegionExtractor initialized without region config - call load_config() before use") | |
| def load_config(self, config_path: str) -> None: | |
| """ | |
| Load play clock region configuration from a JSON file. | |
| Args: | |
| config_path: Path to the JSON config file | |
| """ | |
| path = Path(config_path) | |
| if not path.exists(): | |
| raise FileNotFoundError(f"Config file not found: {config_path}") | |
| with open(path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| self.config = PlayClockRegionConfig( | |
| x_offset=data["x_offset"], | |
| y_offset=data["y_offset"], | |
| width=data["width"], | |
| height=data["height"], | |
| source_video=data.get("source_video", ""), | |
| scorebug_template=data.get("scorebug_template", ""), | |
| samples_used=data.get("samples_used", 0), | |
| ) | |
| logger.info( | |
| "Loaded play clock region config: offset=(%d, %d), size=(%d, %d)", | |
| self.config.x_offset, | |
| self.config.y_offset, | |
| self.config.width, | |
| self.config.height, | |
| ) | |
| def extract_region(self, frame: np.ndarray[Any, Any], scorebug_bbox: Tuple[int, int, int, int]) -> Optional[np.ndarray[Any, Any]]: | |
| """ | |
| Extract the play clock region from the frame. | |
| Args: | |
| frame: Full video frame | |
| scorebug_bbox: Scorebug bounding box (x, y, w, h) | |
| Returns: | |
| Extracted play clock region or None if out of bounds | |
| """ | |
| if self.config is None: | |
| logger.error("No region config loaded - cannot extract play clock region") | |
| return None | |
| sb_x, sb_y, _, _ = scorebug_bbox | |
| # Calculate absolute coordinates of play clock region | |
| pc_x = sb_x + self.config.x_offset | |
| pc_y = sb_y + self.config.y_offset | |
| pc_w = self.config.width | |
| pc_h = self.config.height | |
| # Validate bounds | |
| frame_h, frame_w = frame.shape[:2] | |
| if pc_x < 0 or pc_y < 0 or pc_x + pc_w > frame_w or pc_y + pc_h > frame_h: | |
| logger.warning( | |
| "Play clock region out of bounds: (%d, %d, %d, %d) in frame (%d, %d)", | |
| pc_x, | |
| pc_y, | |
| pc_w, | |
| pc_h, | |
| frame_w, | |
| frame_h, | |
| ) | |
| return None | |
| # Extract region | |
| region = frame[pc_y : pc_y + pc_h, pc_x : pc_x + pc_w].copy() | |
| return region | |
| def preprocess_for_ocr(self, region: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]: | |
| """ | |
| Preprocess the play clock region for OCR (used during template building). | |
| Preprocessing steps: | |
| 1. Detect if digits are red (play clock at 5 seconds or less) | |
| 2. If red, use red channel directly; otherwise convert to grayscale | |
| 3. Scale up for better digit recognition | |
| 4. Apply thresholding (Otsu's or adaptive based on color) | |
| 5. Invert to get dark text on light background | |
| 6. Apply morphological operations to clean up noise | |
| Args: | |
| region: Play clock region (BGR format) | |
| Returns: | |
| Preprocessed image ready for OCR | |
| """ | |
| # Check if digits are red (play clock at 5 seconds or less) | |
| is_red = detect_red_digits(region) | |
| if is_red: | |
| # For red digits, use the red channel directly | |
| _, _, r = cv2.split(region) | |
| gray = r | |
| logger.debug("Using red channel for preprocessing (red play clock detected)") | |
| else: | |
| # Standard grayscale conversion for white digits | |
| gray = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY) | |
| # Scale up by 4x for better OCR accuracy on small digits | |
| scale_factor = 4 | |
| scaled = cv2.resize(gray, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR) | |
| if is_red: | |
| # For red digits, use percentile-based threshold on the red channel | |
| threshold_value = float(np.percentile(np.asarray(scaled), 90)) | |
| _, binary = cv2.threshold(scaled, threshold_value, 255, cv2.THRESH_BINARY) | |
| logger.debug("Red digit threshold (90th percentile): %.1f", threshold_value) | |
| # Apply morphological close to connect digit segments before inverting | |
| kernel = np.ones((2, 2), np.uint8) | |
| binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel) | |
| # Invert to get dark text on light background | |
| binary = cv2.bitwise_not(binary) | |
| else: | |
| # Use Otsu's thresholding for white digits | |
| _, binary = cv2.threshold(scaled, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| # Determine if we need to invert | |
| mean_intensity = np.mean(np.asarray(binary)) | |
| if mean_intensity < 128: | |
| binary = cv2.bitwise_not(binary) | |
| # Apply morphological operations to clean up noise | |
| kernel = np.ones((2, 2), np.uint8) | |
| binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel) | |
| binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel) | |
| # Add padding around the image | |
| padding = 10 | |
| binary = cv2.copyMakeBorder(binary, padding, padding, padding, padding, cv2.BORDER_CONSTANT, value=255) | |
| return binary | |
| def get_absolute_coords(self, scorebug_bbox: Tuple[int, int, int, int]) -> Optional[Tuple[int, int, int, int]]: | |
| """ | |
| Get absolute coordinates of the play clock region given scorebug position. | |
| Args: | |
| scorebug_bbox: Scorebug bounding box (x, y, w, h) | |
| Returns: | |
| Play clock absolute coordinates (x, y, w, h) or None if no config | |
| """ | |
| if self.config is None: | |
| return None | |
| sb_x, sb_y, _, _ = scorebug_bbox | |
| return ( | |
| sb_x + self.config.x_offset, | |
| sb_y + self.config.y_offset, | |
| self.config.width, | |
| self.config.height, | |
| ) | |