evaluate.py

"""
Model evaluation + benchmarking for the custom football YOLO detector.

Produces:
  - mAP@0.5, mAP@0.5:0.95, precision, recall (overall + per class) on the
    held-out test and validation splits, via ultralytics `model.val()`.
  - End-to-end detection throughput (FPS) on the sample match video.
  - A qualitative baseline contrast against the stock COCO yolov8n model,
    which cannot distinguish referees/players and fires on sideline people.
  - Auto-generated plots (PR curve, confusion matrix) under runs/detect/val*/.

Results are written to EVALUATION.md. Re-runnable and side-effect free apart
from that file and the ultralytics run directories.

Usage:
    python evaluate.py
"""

import glob
import os
import subprocess
import time

import cv2
import torch
import yaml
from ultralytics import YOLO

MODEL_PATH = "models/best.pt"
BASELINE_PATH = "yolov8n.pt"
DATA_YAML = "training/data.eval.yaml"
VIDEO_PATH = "Input videos/08fd33_4.mp4"
OUT_MD = "EVALUATION.md"


def pick_device():
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def device_label(device):
    """Human-readable, machine-agnostic device name for the report."""
    if device == "mps":
        try:
            chip = subprocess.check_output(
                ["sysctl", "-n", "machdep.cpu.brand_string"], text=True).strip()
        except Exception:
            chip = "Apple Silicon"
        return f"{chip} (MPS)"
    if device == "cuda":
        try:
            return f"{torch.cuda.get_device_name(0)} (CUDA)"
        except Exception:
            return "NVIDIA GPU (CUDA)"
    return "CPU"


def per_class_rows(metrics):
    """Return list of (class_name, precision, recall, mAP50, mAP50-95)."""
    box = metrics.box
    rows = []
    for i, c in enumerate(box.ap_class_index):
        rows.append((
            metrics.names[c],
            float(box.p[i]),
            float(box.r[i]),
            float(box.ap50[i]),
            float(box.ap[i]),
        ))
    return rows


def count_images(split):
    with open(DATA_YAML) as f:
        cfg = yaml.safe_load(f)
    img_dir = os.path.join(cfg["path"], cfg[split])
    return len(glob.glob(os.path.join(img_dir, "*.*")))


def eval_split(model, split, device):
    print(f"\n=== Evaluating '{split}' split ===")
    m = model.val(data=DATA_YAML, split=split, device=device,
                  plots=True, verbose=False)
    return {
        "split": split,
        "images": count_images(split),
        "map50": float(m.box.map50),
        "map": float(m.box.map),
        "mp": float(m.box.mp),
        "mr": float(m.box.mr),
        "per_class": per_class_rows(m),
        "save_dir": os.path.relpath(str(m.save_dir)),
    }


def measure_fps(model, device, batch_size=20, conf=0.1):
    print("\n=== Measuring detection FPS on sample video ===")
    cap = cv2.VideoCapture(VIDEO_PATH)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)
    cap.release()
    n = len(frames)

    # Warm-up (kernel compilation / lazy init) — excluded from timing.
    model.predict(frames[:batch_size], conf=conf, device=device, verbose=False)

    t0 = time.perf_counter()
    for i in range(0, n, batch_size):
        model.predict(frames[i:i + batch_size], conf=conf,
                      device=device, verbose=False)
    elapsed = time.perf_counter() - t0
    return {"frames": n, "seconds": elapsed, "fps": n / elapsed}


def baseline_contrast(device, conf=0.25):
    """Qualitative: what stock COCO yolov8n does on the same footage."""
    print("\n=== Baseline (stock yolov8n) qualitative contrast ===")
    base = YOLO(BASELINE_PATH)
    cap = cv2.VideoCapture(VIDEO_PATH)
    ret, frame = cap.read()
    cap.release()
    res = base.predict(frame, conf=conf, device=device, verbose=False)[0]
    counts = {}
    for c in res.boxes.cls.tolist():
        name = base.names[int(c)]
        counts[name] = counts.get(name, 0) + 1
    return {"num_classes": len(base.names), "frame0_detections": counts}


def fmt_pct(x):
    return f"{x * 100:.1f}%"


def write_report(device, splits, fps, baseline):
    lines = []
    lines.append("# Model Evaluation\n")
    lines.append(f"- **Detector:** `{MODEL_PATH}` (custom YOLO, 4 classes: "
                 "ball, goalkeeper, player, referee)\n")
    lines.append(f"- **Device:** {device_label(device)}\n")
    lines.append(f"- **Eval config:** `{DATA_YAML}`\n")

    lines.append("\n## Detection accuracy\n")
    lines.append("| Split | Images | mAP@0.5 | mAP@0.5:0.95 | Precision | Recall |\n")
    lines.append("|---|---|---|---|---|---|\n")
    for s in splits:
        lines.append(
            f"| {s['split']} | {s['images']} | {fmt_pct(s['map50'])} "
            f"| {fmt_pct(s['map'])} | {fmt_pct(s['mp'])} | {fmt_pct(s['mr'])} |\n"
        )

    for s in splits:
        lines.append(f"\n### Per-class — `{s['split']}` split\n")
        lines.append("| Class | Precision | Recall | mAP@0.5 | mAP@0.5:0.95 |\n")
        lines.append("|---|---|---|---|---|\n")
        for name, p, r, ap50, ap in s["per_class"]:
            lines.append(f"| {name} | {fmt_pct(p)} | {fmt_pct(r)} "
                         f"| {fmt_pct(ap50)} | {fmt_pct(ap)} |\n")
        lines.append(f"\nPlots (PR curve, confusion matrix): `{s['save_dir']}/`\n")

    lines.append("\n## Inference throughput\n")
    lines.append(f"- Detected **{fps['frames']} frames** in "
                 f"**{fps['seconds']:.1f}s** → **{fps['fps']:.1f} FPS** "
                 f"(detection only, batch 20, on {device_label(device)}).\n")

    lines.append("\n## Baseline contrast — stock COCO `yolov8n`\n")
    lines.append(f"Stock yolov8n has **{baseline['num_classes']} generic COCO "
                 "classes** and no football roles. On frame 0 it produced: "
                 f"`{baseline['frame0_detections']}` — i.e. it collapses "
                 "referees, goalkeepers, and outfield players into a single "
                 "`person` class and offers no team/role signal, motivating the "
                 "custom-trained detector.\n")

    with open(OUT_MD, "w") as f:
        f.writelines(lines)
    print(f"\nWrote {OUT_MD}")


def main():
    device = pick_device()
    print(f"Device: {device}")
    model = YOLO(MODEL_PATH)

    splits = [eval_split(model, "test", device),
              eval_split(model, "val", device)]
    fps = measure_fps(model, device)
    baseline = baseline_contrast(device)
    write_report(device, splits, fps, baseline)


if __name__ == "__main__":
    main()