scripts/archive/v3/configure_timeout_tracker.py

#!/usr/bin/env python3
"""
Interactive tool to configure timeout tracker regions on the scorebug.

This script extracts frames from a video, detects the scorebug, and allows the user
to manually select the timeout indicator regions for each team. The selected regions
are saved to a config file for use by the TimeoutTracker.

Each team has 3 timeout indicators displayed as ovals (white = available, dark = used).

Usage:
    python scripts/configure_timeout_tracker.py

Controls:
    - Click and drag to select a region
    - Press 'h' to save as HOME team timeout region
    - Press 'a' to save as AWAY team timeout region
    - Press 'n' to skip to next frame
    - Press 'r' to reset selection on current frame
    - Press 't' to test detection on current frame
    - Press 'q' to quit and save final config
"""

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

import cv2
import numpy as np

from detection import DetectScoreBug, TrackTimeouts, TimeoutRegionConfig
from ui import RegionSelector

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

# Constants
VIDEO_PATH = Path(__file__).parent.parent / "full_videos" / "OSU vs Tenn 12.21.24.mkv"
TEMPLATE_PATH = Path(__file__).parent.parent / "data" / "templates" / "scorebug_template_main.png"
CONFIG_OUTPUT_PATH = Path(__file__).parent.parent / "data" / "config" / "timeout_tracker_region.json"

# Test segment: beginning of game when both teams have all 3 timeouts
# Avoid times when graphics like "IMPACT PLAYERS" might be showing
START_TIME_SECONDS = 2 * 60 + 40  # 2:40 - early in first quarter after initial graphics
END_TIME_SECONDS = 4 * 60 + 30  # 4:30
SAMPLE_INTERVAL_SECONDS = 3  # Sample every 3 seconds


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(f"Could not open video: {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 - require high confidence to avoid picking up graphics overlays
        detection = detector.detect(frame)
        if detection.detected and detection.bbox and detection.confidence >= 0.90:
            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 (conf=%.2f)", current_time, detection.confidence if detection else 0.0)

        current_time += interval

    cap.release()
    return frames_with_scorebug


def convert_to_absolute_bbox(relative_bbox: Tuple[int, int, int, int], scorebug_bbox: Tuple[int, int, int, int], scale_factor: int) -> Tuple[int, int, int, int]:
    """Convert a selection made on scaled display to absolute frame coordinates."""
    sb_x, sb_y, _, _ = scorebug_bbox
    rel_x, rel_y, rel_w, rel_h = relative_bbox

    # Convert from display scale to original scale
    original_rel_x = rel_x // scale_factor
    original_rel_y = rel_y // scale_factor
    original_rel_w = rel_w // scale_factor
    original_rel_h = rel_h // scale_factor

    # Convert to absolute frame coordinates
    abs_x = sb_x + original_rel_x
    abs_y = sb_y + original_rel_y

    return (abs_x, abs_y, original_rel_w, original_rel_h)


def test_timeout_detection(frame: np.ndarray, tracker: TrackTimeouts) -> None:
    """Test and display timeout detection on a frame."""
    if not tracker.is_configured():
        logger.warning("Tracker not fully configured yet")
        return

    reading = tracker.read_timeouts(frame)
    logger.info("Test detection:")
    logger.info("  Home timeouts: %d (states: %s)", reading.home_timeouts, reading.home_oval_states)
    logger.info("  Away timeouts: %d (states: %s)", reading.away_timeouts, reading.away_oval_states)
    logger.info("  Confidence: %.2f", reading.confidence)

    # Show visualization
    vis_frame = tracker.visualize(frame, reading)
    vis_frame = cv2.resize(vis_frame, (vis_frame.shape[1] // 2, vis_frame.shape[0] // 2))
    cv2.imshow("Timeout Detection Test", vis_frame)
    cv2.waitKey(2000)  # Show for 2 seconds
    cv2.destroyWindow("Timeout Detection Test")


def run_region_selection(frames_with_scorebug: List[Tuple[float, Any, Tuple[int, int, int, int]]]) -> Tuple[Optional[TimeoutRegionConfig], Optional[TimeoutRegionConfig]]:
    """
    Run interactive region selection for timeout indicators.

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

    Returns:
        Tuple of (home_region_config, away_region_config), either may be None
    """
    if not frames_with_scorebug:
        logger.error("No frames with scorebug detected!")
        return None, None

    window_name = "Select Timeout Regions (h=home, a=away)"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

    # Use drag mode for this script's existing behavior
    selector = RegionSelector(window_name, mode="drag")
    cv2.setMouseCallback(window_name, selector.mouse_callback)

    home_region: Optional[TimeoutRegionConfig] = None
    away_region: Optional[TimeoutRegionConfig] = None

    # Create tracker for testing
    tracker = TrackTimeouts()

    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 25 pixels in original scale)
            grid_spacing = 25 * 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 (green while selecting)
            if selector.start_point and selector.end_point:
                cv2.rectangle(display_with_selection, selector.start_point, selector.end_point, (0, 255, 0), 2)

            # Draw already-saved regions
            if home_region:
                # Convert absolute coords to display coords
                rel_x = (home_region.bbox[0] - sb_x) * scale_factor
                rel_y = (home_region.bbox[1] - sb_y) * scale_factor
                rel_w = home_region.bbox[2] * scale_factor
                rel_h = home_region.bbox[3] * scale_factor
                cv2.rectangle(display_with_selection, (rel_x, rel_y), (rel_x + rel_w, rel_y + rel_h), (255, 0, 0), 2)  # Blue for home
                cv2.putText(display_with_selection, "HOME", (rel_x, rel_y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)

            if away_region:
                rel_x = (away_region.bbox[0] - sb_x) * scale_factor
                rel_y = (away_region.bbox[1] - sb_y) * scale_factor
                rel_w = away_region.bbox[2] * scale_factor
                rel_h = away_region.bbox[3] * scale_factor
                cv2.rectangle(display_with_selection, (rel_x, rel_y), (rel_x + rel_w, rel_y + rel_h), (0, 165, 255), 2)  # Orange for away
                cv2.putText(display_with_selection, "AWAY", (rel_x, rel_y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 165, 255), 1)

            # Add instructions text
            instructions = [
                f"Frame {frame_idx + 1}/{len(frames_with_scorebug)} @ {timestamp:.1f}s",
                "Select 3-oval timeout region for each team",
                "h=save HOME, a=save AWAY, n=next, r=reset",
                "t=test detection, q=quit and save",
                f"HOME: {'SET' if home_region else 'NOT SET'}  AWAY: {'SET' if away_region else 'NOT SET'}",
            ]
            y_pos = 30
            for text in instructions:
                cv2.putText(display_with_selection, text, (10, y_pos), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
                y_pos += 20

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

            # Handle key presses
            if key == ord("q"):
                cv2.destroyAllWindows()
                return home_region, away_region

            if key == ord("h"):
                # Save as HOME team region
                bbox = selector.get_bbox()
                if bbox:
                    abs_bbox = convert_to_absolute_bbox(bbox, scorebug_bbox, scale_factor)
                    home_region = TimeoutRegionConfig(team_name="home", bbox=abs_bbox)
                    logger.info("Saved HOME region: %s", abs_bbox)
                    # Update tracker for testing
                    if away_region:
                        tracker.set_regions(home_region, away_region)
                selector.reset()

            elif key == ord("a"):
                # Save as AWAY team region
                bbox = selector.get_bbox()
                if bbox:
                    abs_bbox = convert_to_absolute_bbox(bbox, scorebug_bbox, scale_factor)
                    away_region = TimeoutRegionConfig(team_name="away", bbox=abs_bbox)
                    logger.info("Saved AWAY region: %s", abs_bbox)
                    # Update tracker for testing
                    if home_region:
                        tracker.set_regions(home_region, away_region)
                selector.reset()

            elif key == ord("n"):
                # Skip to next frame
                selector.reset()
                frame_idx += 1
                break

            elif key == ord("r"):
                # Reset selection
                selector.reset()

            elif key == ord("t"):
                # Test detection
                if home_region and away_region:
                    tracker.set_regions(home_region, away_region)
                    test_timeout_detection(frame, tracker)
                else:
                    logger.warning("Configure both regions before testing")

    cv2.destroyAllWindows()
    return home_region, away_region


def save_config(home_region: TimeoutRegionConfig, away_region: TimeoutRegionConfig, output_path: Path) -> None:
    """Save the timeout tracker configuration to a JSON file."""
    data = {
        "home_timeout_region": home_region.to_dict(),
        "away_timeout_region": away_region.to_dict(),
        "source_video": str(VIDEO_PATH.name),
        "scorebug_template": str(TEMPLATE_PATH.name),
    }

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

    logger.info("Saved config to %s", output_path)


def main():
    """Main entry point for timeout tracker region configuration."""
    logger.info("Timeout Tracker Configuration 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 a 3-oval timeout region")
    logger.info("  - Press 'h' to save selection as HOME team region")
    logger.info("  - Press 'a' to save selection as AWAY team region")
    logger.info("  - Press 'n' to skip to next frame")
    logger.info("  - Press 'r' to reset current selection")
    logger.info("  - Press 't' to test detection with current config")
    logger.info("  - Press 'q' to quit and save")

    home_region, away_region = run_region_selection(frames)

    if home_region and away_region:
        logger.info("Configuration complete!")
        logger.info("  HOME region: %s", home_region.bbox)
        logger.info("  AWAY region: %s", away_region.bbox)
        save_config(home_region, away_region, CONFIG_OUTPUT_PATH)
        return 0

    if home_region or away_region:
        logger.warning("Only one region configured. Both are required.")
        if home_region:
            logger.info("  HOME region: %s", home_region.bbox)
        if away_region:
            logger.info("  AWAY region: %s", away_region.bbox)
        # Save partial config anyway
        if home_region and away_region:
            save_config(home_region, away_region, CONFIG_OUTPUT_PATH)
        return 1

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


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