scripts/archive/v2/identify_play_clock_region.py

#!/usr/bin/env python3
"""
Interactive tool to identify the play clock region within the scorebug.

This script extracts frames from a video, detects the scorebug using template matching,
and allows the user to manually select the play clock sub-region within the scorebug.
The selected region coordinates are saved to a config file for use by the PlayClockReader.

Usage:
    python scripts/identify_play_clock_region.py

Controls:
    - Click and drag to select the play clock region
    - Press 's' to save the selection and move to next frame
    - Press 'n' to skip to next frame without saving
    - Press 'r' to reset selection on current frame
    - Press 'q' to quit and save final config
"""

import json
import logging
import sys
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Optional, Tuple, List, Any

import cv2

from detection import DetectScoreBug

# Path reference for constants
PROJECT_ROOT = Path(__file__).parent.parent.parent

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Constants
VIDEO_PATH = PROJECT_ROOT / "full_videos" / "OSU vs Tenn 12.21.24.mkv"
TEMPLATE_PATH = PROJECT_ROOT / "data" / "templates" / "scorebug_template_main.png"
CONFIG_OUTPUT_PATH = PROJECT_ROOT / "data" / "config" / "play_clock_region.json"

# Test segment: 38:40 to 39:00 (20 seconds for sampling frames)
START_TIME_SECONDS = 38 * 60 + 40  # 38:40
END_TIME_SECONDS = 39 * 60  # 39:00
SAMPLE_INTERVAL_SECONDS = 2  # Sample every 2 seconds


@dataclass
class PlayClockRegionConfig:
    """Configuration for the play clock region relative to the scorebug bounding box."""

    # Relative offsets from scorebug bbox origin (top-left)
    x_offset: int  # X offset from scorebug left edge
    y_offset: int  # Y offset from scorebug top edge
    width: int  # Width of play clock region
    height: int  # Height of play clock region

    # Metadata
    source_video: str  # Video used to identify region
    scorebug_template: str  # Template used for scorebug detection
    samples_used: int  # Number of frames used to verify region


class RegionSelector:
    """Interactive region selector using OpenCV mouse callbacks."""

    def __init__(self, window_name: str):
        self.window_name = window_name
        self.start_point: Optional[Tuple[int, int]] = None
        self.end_point: Optional[Tuple[int, int]] = None
        self.drawing = False
        self.selection_complete = False

    def mouse_callback(self, event, x, y, flags, param):  # pylint: disable=unused-argument
        """Handle mouse events for region selection."""
        if event == cv2.EVENT_LBUTTONDOWN:
            # Start drawing rectangle
            self.start_point = (x, y)
            self.end_point = (x, y)
            self.drawing = True
            self.selection_complete = False
        elif event == cv2.EVENT_MOUSEMOVE and self.drawing:
            # Update rectangle as mouse moves
            self.end_point = (x, y)
        elif event == cv2.EVENT_LBUTTONUP:
            # Finish drawing rectangle
            self.end_point = (x, y)
            self.drawing = False
            self.selection_complete = True

    def get_bbox(self) -> Optional[Tuple[int, int, int, int]]:
        """Get the selected bounding box as (x, y, width, height)."""
        if self.start_point is None or self.end_point is None:
            return None

        x1, y1 = self.start_point
        x2, y2 = self.end_point

        # Normalize to ensure positive width/height
        x = min(x1, x2)
        y = min(y1, y2)
        w = abs(x2 - x1)
        h = abs(y2 - y1)

        if w == 0 or h == 0:
            return None

        return (x, y, w, h)

    def reset(self):
        """Reset the selection."""
        self.start_point = None
        self.end_point = None
        self.drawing = False
        self.selection_complete = False


def extract_frames_with_scorebug(
    video_path: Path, detector: DetectScoreBug, start_time: float, end_time: float, interval: float
) -> List[Tuple[float, Any, Tuple[int, int, int, int]]]:
    """
    Extract frames from video where scorebug is detected.

    Args:
        video_path: Path to video file
        detector: DetectScoreBug instance
        start_time: Start time in seconds
        end_time: End time in seconds
        interval: Sampling interval in seconds

    Returns:
        List of (timestamp, frame, scorebug_bbox) tuples
    """
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise ValueError("Could not open video: %s" % video_path)

    fps = cap.get(cv2.CAP_PROP_FPS)
    logger.info("Video FPS: %s", fps)

    frames_with_scorebug = []
    current_time = start_time

    while current_time < end_time:
        # Seek to current time
        frame_number = int(current_time * fps)
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)

        ret, frame = cap.read()
        if not ret:
            logger.warning("Could not read frame at %.1fs", current_time)
            current_time += interval
            continue

        # Detect scorebug
        detection = detector.detect(frame)
        if detection.detected and detection.bbox:
            logger.info("Scorebug detected at %.1fs with confidence %.2f", current_time, detection.confidence)
            frames_with_scorebug.append((current_time, frame, detection.bbox))
        else:
            logger.debug("No scorebug at %.1fs", current_time)

        current_time += interval

    cap.release()
    return frames_with_scorebug


def _handle_key_press(key: int, selector: RegionSelector, selections: List, scale_factor: int) -> Optional[str]:
    """
    Handle keyboard input during region selection.

    Returns:
        'quit' to quit and return, 'next' to advance frame, None to continue loop
    """
    if key == ord("q"):
        return "quit"
    if key == ord("s"):
        # Save selection and move to next frame
        bbox = selector.get_bbox()
        if bbox:
            # Convert from display scale back to original
            x, y, w, h = bbox
            original_bbox = (x // scale_factor, y // scale_factor, w // scale_factor, h // scale_factor)
            selections.append(original_bbox)
            logger.info("Saved selection: %s", original_bbox)
        selector.reset()
        return "next"
    if key == ord("n"):
        # Skip to next frame
        selector.reset()
        return "next"
    if key == ord("r"):
        # Reset selection
        selector.reset()
    return None


def run_region_selection(frames_with_scorebug: List[Tuple[float, Any, Tuple[int, int, int, int]]]) -> Optional[Tuple[int, int, int, int]]:
    """
    Run interactive region selection on frames with scorebug.

    Args:
        frames_with_scorebug: List of (timestamp, frame, scorebug_bbox) tuples

    Returns:
        Selected region as (x_offset, y_offset, width, height) relative to scorebug, or None if cancelled
    """
    if not frames_with_scorebug:
        logger.error("No frames with scorebug detected!")
        return None

    window_name = "Select Play Clock Region (click and drag)"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

    selector = RegionSelector(window_name)
    cv2.setMouseCallback(window_name, selector.mouse_callback)

    # Store all selections to find consensus
    selections: List[Tuple[int, int, int, int]] = []
    frame_idx = 0
    scale_factor = 2  # Scale up for easier selection

    while frame_idx < len(frames_with_scorebug):
        timestamp, frame, scorebug_bbox = frames_with_scorebug[frame_idx]
        sb_x, sb_y, sb_w, sb_h = scorebug_bbox

        # Extract and enlarge scorebug region for easier selection
        scorebug_region = frame[sb_y : sb_y + sb_h, sb_x : sb_x + sb_w].copy()
        display_region = cv2.resize(scorebug_region, (sb_w * scale_factor, sb_h * scale_factor), interpolation=cv2.INTER_LINEAR)

        while True:
            # Draw current selection on display image
            display_with_selection = display_region.copy()

            # Draw grid lines for reference (every 50 pixels in original scale)
            grid_spacing = 50 * scale_factor
            for gx in range(0, display_with_selection.shape[1], grid_spacing):
                cv2.line(display_with_selection, (gx, 0), (gx, display_with_selection.shape[0]), (100, 100, 100), 1)
            for gy in range(0, display_with_selection.shape[0], grid_spacing):
                cv2.line(display_with_selection, (0, gy), (display_with_selection.shape[1], gy), (100, 100, 100), 1)

            # Draw current selection rectangle
            if selector.start_point and selector.end_point:
                cv2.rectangle(display_with_selection, selector.start_point, selector.end_point, (0, 255, 0), 2)

            # Add instructions text
            instructions = [
                "Frame %d/%d @ %.1fs" % (frame_idx + 1, len(frames_with_scorebug), timestamp),
                "Click and drag to select play clock",
                "s=save, n=next, r=reset, q=quit",
            ]
            y_pos = 30
            for text in instructions:
                cv2.putText(display_with_selection, text, (10, y_pos), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
                y_pos += 25

            cv2.imshow(window_name, display_with_selection)
            key = cv2.waitKey(30) & 0xFF

            action = _handle_key_press(key, selector, selections, scale_factor)
            if action == "quit":
                cv2.destroyAllWindows()
                if selections:
                    return _average_selections(selections)
                return None
            if action == "next":
                frame_idx += 1
                break

    cv2.destroyAllWindows()

    if selections:
        return _average_selections(selections)
    return None


def _average_selections(selections: List[Tuple[int, int, int, int]]) -> Tuple[int, int, int, int]:
    """Average multiple selections to get consensus region."""
    if not selections:
        return (0, 0, 0, 0)

    avg_x = sum(s[0] for s in selections) // len(selections)
    avg_y = sum(s[1] for s in selections) // len(selections)
    avg_w = sum(s[2] for s in selections) // len(selections)
    avg_h = sum(s[3] for s in selections) // len(selections)

    return (avg_x, avg_y, avg_w, avg_h)


def save_config(region: Tuple[int, int, int, int], output_path: Path, samples_used: int):
    """Save the play clock region configuration to a JSON file."""
    config = PlayClockRegionConfig(
        x_offset=region[0],
        y_offset=region[1],
        width=region[2],
        height=region[3],
        source_video=str(VIDEO_PATH.name),
        scorebug_template=str(TEMPLATE_PATH.name),
        samples_used=samples_used,
    )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(asdict(config), f, indent=2)

    logger.info("Saved config to %s", output_path)
    logger.info("Play clock region: x_offset=%d, y_offset=%d, width=%d, height=%d", config.x_offset, config.y_offset, config.width, config.height)


def main():
    """Main entry point for play clock region identification."""
    logger.info("Play Clock Region Identification Tool")
    logger.info("=" * 50)

    # Verify paths exist
    if not VIDEO_PATH.exists():
        logger.error("Video not found: %s", VIDEO_PATH)
        logger.info("Expected video at: full_videos/OSU vs Tenn 12.21.24.mkv")
        return 1

    if not TEMPLATE_PATH.exists():
        logger.error("Template not found: %s", TEMPLATE_PATH)
        return 1

    # Initialize scorebug detector
    logger.info("Loading scorebug template: %s", TEMPLATE_PATH)
    detector = DetectScoreBug(template_path=str(TEMPLATE_PATH))

    # Extract frames with scorebug
    logger.info("Extracting frames from %ds to %ds...", START_TIME_SECONDS, END_TIME_SECONDS)
    frames = extract_frames_with_scorebug(VIDEO_PATH, detector, START_TIME_SECONDS, END_TIME_SECONDS, SAMPLE_INTERVAL_SECONDS)

    if not frames:
        logger.error("No frames with scorebug detected! Check template matching.")
        return 1

    logger.info("Found %d frames with scorebug", len(frames))

    # Run interactive selection
    logger.info("Starting interactive selection...")
    logger.info("Instructions:")
    logger.info("  - Click and drag to select the play clock region")
    logger.info("  - Press 's' to save selection and move to next frame")
    logger.info("  - Press 'n' to skip to next frame")
    logger.info("  - Press 'r' to reset selection")
    logger.info("  - Press 'q' to quit and save")

    selected_region = run_region_selection(frames)

    if selected_region:
        logger.info("Final selected region: %s", selected_region)
        save_config(selected_region, CONFIG_OUTPUT_PATH, len(frames))
        logger.info("Region identification complete!")
        return 0

    logger.warning("No region selected. Exiting without saving.")
    return 1


if __name__ == "__main__":
    sys.exit(main())