debug_ball_detect.py

"""Diagnostic tool for ball detection.

Writes annotated frames to debug_balls/ showing:
- HSV yellow mask
- All yellow blob candidates (cyan circles)
- Launch zone (yellow rect)
- Robot bbox (green rect)

Usage:
  uv run python debug_ball_detect.py saves/q2.saproj
  uv run python debug_ball_detect.py saves/q2.saproj --start-s 30 --duration-s 5
  uv run python debug_ball_detect.py saves/q2.saproj --h-low 15 --s-low 50

Everything defaults from the .saproj file.
"""

from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

import cv2
import numpy as np

from gui.project import load_project
from gui.ball_detector import (
    _robot_bbox_at,
    _filter_contours,
    _in_launch_zone,
)
from utils.config import BallDetectionConfig


def main():
    p = argparse.ArgumentParser(
        description="Debug ball detection visually",
    )
    p.add_argument(
        "project", help=".saproj with tracking data",
    )
    p.add_argument(
        "--video", default=None,
        help="Override video path (default: from project)",
    )
    p.add_argument(
        "--start-s", type=float, default=None,
        help="Start time in seconds "
        "(default: first tracked frame)",
    )
    p.add_argument(
        "--duration-s", type=float, default=10,
        help="Duration to analyze (default: 10s)",
    )
    p.add_argument(
        "--every-n", type=int, default=3,
        help="Save every Nth frame (default: 3)",
    )
    p.add_argument(
        "--out", default="debug_balls",
        help="Output directory (default: debug_balls)",
    )
    # HSV overrides
    p.add_argument("--h-low", type=int, default=None)
    p.add_argument("--h-high", type=int, default=None)
    p.add_argument("--s-low", type=int, default=None)
    p.add_argument("--s-high", type=int, default=None)
    p.add_argument("--v-low", type=int, default=None)
    p.add_argument("--v-high", type=int, default=None)
    p.add_argument(
        "--min-area", type=int, default=None,
    )
    p.add_argument(
        "--max-area", type=int, default=None,
    )
    args = p.parse_args()

    # Load project
    proj = load_project(Path(args.project))
    if not proj.tracking_results:
        print("ERROR: No tracking results in project")
        sys.exit(1)

    video = args.video or proj.video_path
    if not video:
        print("ERROR: No video path in project "
              "(use --video)")
        sys.exit(1)

    tracking = proj.tracking_results
    track_times = [r.time_ms for r in tracking]

    cfg = (
        proj.ball_detection_config
        or BallDetectionConfig()
    )

    # Apply CLI overrides
    if args.h_low is not None:
        cfg.ball_hsv.h_low = args.h_low
    if args.h_high is not None:
        cfg.ball_hsv.h_high = args.h_high
    if args.s_low is not None:
        cfg.ball_hsv.s_low = args.s_low
    if args.s_high is not None:
        cfg.ball_hsv.s_high = args.s_high
    if args.v_low is not None:
        cfg.ball_hsv.v_low = args.v_low
    if args.v_high is not None:
        cfg.ball_hsv.v_high = args.v_high
    if args.min_area is not None:
        cfg.min_ball_area = args.min_area
    if args.max_area is not None:
        cfg.max_ball_area = args.max_area

    hsv = cfg.ball_hsv
    lower = np.array(
        [hsv.h_low, hsv.s_low, hsv.v_low],
        dtype=np.uint8,
    )
    upper = np.array(
        [hsv.h_high, hsv.s_high, hsv.v_high],
        dtype=np.uint8,
    )
    kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (cfg.morph_kernel_size,
         cfg.morph_kernel_size),
    )

    # Default start to first tracked frame
    if args.start_s is not None:
        start_ms = args.start_s * 1000
    else:
        start_ms = track_times[0] if track_times else 0

    end_ms = start_ms + args.duration_s * 1000

    print(f"Video: {video}")
    print(f"Scanning: {start_ms / 1000:.1f}s — "
          f"{end_ms / 1000:.1f}s")
    print(f"HSV range: H[{hsv.h_low}-{hsv.h_high}] "
          f"S[{hsv.s_low}-{hsv.s_high}] "
          f"V[{hsv.v_low}-{hsv.v_high}]")
    print(f"Area range: [{cfg.min_ball_area}, "
          f"{cfg.max_ball_area}]")
    print(f"Launch zone: height_mult="
          f"{cfg.launch_zone_height_mult}, "
          f"width_pad={cfg.launch_zone_width_pad}")
    print()

    # Prepare output dir
    out_dir = Path(args.out)
    if out_dir.exists():
        shutil.rmtree(out_dir)
    out_dir.mkdir(parents=True)

    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        print(f"ERROR: Cannot open {video}")
        sys.exit(1)

    cap.set(cv2.CAP_PROP_POS_MSEC, start_ms)
    frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

    frame_idx = 0
    saved = 0
    total_candidates = 0
    total_in_lz = 0
    frames_with_candidates = 0
    frames_no_robot = 0

    print(f"{'frame':>6} {'time_s':>8} {'robot':>6} "
          f"{'cands':>6} {'in_lz':>6} {'areas'}")
    print("-" * 60)

    while True:
        ok, frame = cap.read()
        if not ok:
            break
        t_ms = cap.get(cv2.CAP_PROP_POS_MSEC)
        if t_ms > end_ms:
            break

        robot_bbox = _robot_bbox_at(
            t_ms, tracking, track_times,
        )

        # HSV mask
        hsv_frame = cv2.cvtColor(
            frame, cv2.COLOR_BGR2HSV,
        )
        mask = cv2.inRange(hsv_frame, lower, upper)
        mask_clean = cv2.morphologyEx(
            mask, cv2.MORPH_OPEN, kernel,
        )
        mask_clean = cv2.morphologyEx(
            mask_clean, cv2.MORPH_CLOSE, kernel,
        )

        # All contours (no area filter)
        contours_raw, _ = cv2.findContours(
            mask_clean,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE,
        )

        # Filtered candidates
        candidates = _filter_contours(
            contours_raw,
            cfg.min_ball_area,
            cfg.max_ball_area,
        )

        # Launch zone
        lz = None
        in_lz_count = 0
        if robot_bbox is not None:
            rx, ry, rw, rh = robot_bbox
            pad_x = int(
                rw * cfg.launch_zone_width_pad,
            )
            lz_x1 = max(0, rx - pad_x)
            lz_x2 = min(frame_w, rx + rw + pad_x)
            lz_y1 = max(
                0,
                ry - int(
                    rh * cfg.launch_zone_height_mult
                ),
            )
            lz_y2 = ry + int(rh * 0.25)
            lz = (lz_x1, lz_y1, lz_x2, lz_y2)

            for cx, cy, _ in candidates:
                if _in_launch_zone(
                    cx, cy,
                    lz_x1, lz_y1, lz_x2, lz_y2,
                    pad_x, rh,
                ):
                    in_lz_count += 1
        else:
            frames_no_robot += 1

        if candidates:
            frames_with_candidates += 1
        total_candidates += len(candidates)
        total_in_lz += in_lz_count

        # Top 5 contour areas for debugging
        all_areas = sorted(
            [cv2.contourArea(c) for c in contours_raw],
            reverse=True,
        )[:5]

        # Print stats
        has_robot = "Y" if robot_bbox else "N"
        areas_str = ", ".join(
            f"{a:.0f}" for a in all_areas
        )
        print(
            f"{frame_idx:>6} {t_ms / 1000:>8.2f} "
            f"{has_robot:>6} {len(candidates):>6} "
            f"{in_lz_count:>6} [{areas_str}]"
        )

        # Save annotated frame every N frames
        if frame_idx % args.every_n == 0:
            annotated = frame.copy()

            # Robot bbox (green)
            if robot_bbox is not None:
                rx, ry, rw, rh = robot_bbox
                cv2.rectangle(
                    annotated,
                    (rx, ry), (rx + rw, ry + rh),
                    (0, 255, 0), 2,
                )

            # Launch zone (yellow)
            if lz is not None:
                cv2.rectangle(
                    annotated,
                    (lz[0], lz[1]), (lz[2], lz[3]),
                    (0, 255, 255), 2,
                )
                cv2.putText(
                    annotated, "LZ",
                    (lz[0] + 2, lz[1] + 15),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    (0, 255, 255), 1,
                )

            # Draw all contours with color coding
            for cnt in contours_raw:
                area = cv2.contourArea(cnt)
                M = cv2.moments(cnt)
                if M["m00"] == 0:
                    continue
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                if area < cfg.min_ball_area:
                    # Too small — red dot
                    cv2.circle(
                        annotated, (cx, cy), 3,
                        (0, 0, 200), -1,
                    )
                elif area > cfg.max_ball_area:
                    # Too big — magenta dot
                    cv2.circle(
                        annotated, (cx, cy), 6,
                        (255, 0, 255), 2,
                    )
                else:
                    # Valid candidate — cyan circle
                    r = int((area / 3.14) ** 0.5)
                    cv2.circle(
                        annotated, (cx, cy), r,
                        (255, 255, 0), 2,
                    )
                    cv2.putText(
                        annotated,
                        f"{area:.0f}",
                        (cx + r + 2, cy),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.4, (255, 255, 0), 1,
                    )

            # Side-by-side: annotated + mask
            mask_bgr = cv2.cvtColor(
                mask_clean, cv2.COLOR_GRAY2BGR,
            )
            combined = np.hstack(
                [annotated, mask_bgr],
            )

            fname = (
                out_dir
                / f"{saved:04d}_{t_ms:.0f}ms.jpg"
            )
            cv2.imwrite(str(fname), combined)
            saved += 1

        frame_idx += 1

    cap.release()

    print()
    print(f"=== Summary ({frame_idx} frames) ===")
    print(f"  Frames with no robot bbox: "
          f"{frames_no_robot}")
    print(f"  Frames with candidates: "
          f"{frames_with_candidates}")
    print(f"  Total candidates: {total_candidates}")
    print(f"  Total in launch zone: {total_in_lz}")
    print(f"  Saved {saved} annotated frames "
          f"to {out_dir}/")
    print()
    print("Color key:")
    print("  Green rect  = robot bbox")
    print("  Yellow rect = launch zone")
    print("  Cyan circle = valid candidate (w/ area)")
    print("  Red dot     = too small (< min_area)")
    print("  Magenta dot = too big (> max_area)")
    print("  Right half  = HSV mask (white = yellow)")


if __name__ == "__main__":
    main()