team_visual_agent.py

"""Visual team evidence for basketball highlight candidates.

This is deliberately conservative: it only labels a candidate when a COCO
person detector finds a likely ball-handler and that player's upper-body jersey
has a clear Lakers-yellow or Celtics-green signal.  Ambiguous frames stay
"unknown" rather than contaminating a team-specific edit.
"""
from __future__ import annotations

from collections import Counter
from pathlib import Path
from typing import Any


_LAKERS = "湖人"
_CELTICS = "凯尔特人"
_UNKNOWN = "待确认"


def _jersey_vote(frame: Any, box: tuple[int, int, int, int]) -> tuple[str, float]:
    """Classify an upper-body crop by strong yellow/green jersey pixels."""
    import cv2
    import numpy as np

    x1, y1, x2, y2 = box
    width, height = max(1, x2 - x1), max(1, y2 - y1)
    # Avoid heads, skin and shorts. The middle torso is the least noisy region.
    crop = frame[y1 + int(height * 0.24): y1 + int(height * 0.68), x1 + int(width * 0.18): x1 + int(width * 0.82)]
    if crop.size == 0:
        return _UNKNOWN, 0.0
    hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV)
    yellow = cv2.inRange(hsv, np.array([15, 75, 80]), np.array([42, 255, 255]))
    green = cv2.inRange(hsv, np.array([38, 55, 35]), np.array([95, 255, 235]))
    yellow_score = float((yellow > 0).mean())
    green_score = float((green > 0).mean())
    if yellow_score >= 0.055 and yellow_score > green_score * 1.3:
        return _LAKERS, min(0.95, yellow_score * 5.0)
    if green_score >= 0.055 and green_score > yellow_score * 1.3:
        return _CELTICS, min(0.95, green_score * 5.0)
    return _UNKNOWN, max(yellow_score, green_score)


def _candidate_times(candidate: dict[str, Any]) -> list[float]:
    start, end = float(candidate["start_sec"]), float(candidate["end_sec"])
    if end - start < 1.0:
        return [start]
    return [start + (end - start) * ratio for ratio in (0.32, 0.55, 0.78)]


def analyze_team_evidence(
    video_path: str | Path,
    candidates: list[dict[str, Any]],
    device: str = "cuda",
) -> tuple[dict[str, dict[str, Any]], str]:
    """Return team evidence keyed by candidate id without modifying clip JSON.

    YOLO detects people and the COCO sports-ball class.  The person nearest the
    ball is used as a likely shooter; their jersey colour determines the vote.
    No model/dependency failure can stop the editing pipeline.
    """
    if not candidates:
        return {}, "没有可供识别的高光候选。"
    try:
        import cv2
        from ultralytics import YOLO
    except ImportError:
        return {}, "视觉模块未安装：请先执行 pip install -r requirements.txt。"

    try:
        # Ultralytics downloads this compact official detector on first use and
        # then reuses the local weight cache. device=0 means the first CUDA GPU.
        model = YOLO("yolo26n.pt")
        cap = cv2.VideoCapture(str(video_path))
        fps = float(cap.get(cv2.CAP_PROP_FPS) or 30.0)
        if not cap.isOpened():
            return {}, "视觉模块无法读取上传视频。"
        evidence: dict[str, dict[str, Any]] = {}
        yolo_device: int | str = 0 if device == "cuda" else "cpu"

        for candidate in candidates:
            votes: list[tuple[str, float]] = []
            checked = 0
            for second in _candidate_times(candidate):
                cap.set(cv2.CAP_PROP_POS_FRAMES, int(second * fps))
                ok, frame = cap.read()
                if not ok:
                    continue
                checked += 1
                result = model(frame, imgsz=640, conf=0.22, classes=[0, 32], device=yolo_device, verbose=False)[0]
                boxes = result.boxes
                if boxes is None:
                    continue
                xyxy = boxes.xyxy.cpu().tolist()
                classes = boxes.cls.cpu().tolist()
                confidences = boxes.conf.cpu().tolist()
                people = [(tuple(map(int, coords)), float(conf)) for coords, cls, conf in zip(xyxy, classes, confidences) if int(cls) == 0]
                balls = [coords for coords, cls in zip(xyxy, classes) if int(cls) == 32]
                if not people or not balls:
                    continue
                ball = max(balls, key=lambda item: (item[2] - item[0]) * (item[3] - item[1]))
                ball_x, ball_y = (ball[0] + ball[2]) / 2, (ball[1] + ball[3]) / 2
                shooter, person_conf = min(
                    people,
                    key=lambda item: ((item[0][0] + item[0][2]) / 2 - ball_x) ** 2 + ((item[0][1] + item[0][3]) / 2 - ball_y) ** 2,
                )
                team, colour_conf = _jersey_vote(frame, shooter)
                if team != _UNKNOWN:
                    votes.append((team, min(0.95, 0.45 + colour_conf * 0.35 + person_conf * 0.2)))

            totals: Counter[str] = Counter()
            for team, confidence in votes:
                totals[team] += confidence
            if totals:
                team, score = totals.most_common(1)[0]
                total = sum(totals.values())
                confidence = round(min(0.95, score / max(total, 0.001) * (0.45 + 0.12 * len(votes))), 2)
                # A close yellow/green split is not trustworthy enough for a
                # team-specific cut.
                if len(totals) > 1 and score / total < 0.68:
                    team, confidence = _UNKNOWN, 0.0
            else:
                team, confidence = _UNKNOWN, 0.0
            evidence[str(candidate["candidate_id"])] = {
                "team": team,
                "confidence": confidence,
                "method": "持球人附近的球衣颜色" if team != _UNKNOWN else "未检测到可确认的持球球衣",
                "frames_checked": checked,
            }
        cap.release()
        return evidence, "视觉队伍识别已完成（黄色=湖人，绿色=凯尔特人）。"
    except Exception as exc:
        return {}, f"视觉队伍识别未完成：{type(exc).__name__}。"