热搜第一！常州队进球无效？用YOLO建个足球追踪系统试试看！

/data
  └── images/
      └── match01_001.jpg
  └── labels/
      └── match01_001.txt  # YOLO格式标签

      yolo task=detect mode=train model=yolov8n.pt data=match.yaml epochs=50 imgsz=640

from ultralytics import YOLO

model = YOLO("../assets/weights/best.pt")

result = model.predict(source="../assets/source_videos/B1606b0e6_1 (28).mp4", 
                       save=True,
                       stream=True)

import cv2

# Retrieve and visualize the first result
first_result = next(result)  # Get the first result from the generator
annotated_frame = first_result.plot()  # Overlay detections on the frame

# Display the annotated frame
cv2.imshow("YOLO Detection", annotated_frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

# funcs.py
import cv2
import numpy as np
from typing import List

def read_video_frames(video_path: str) -> List[np.ndarray]:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise IOError(f"Cannot open video: {video_path}")

    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)
    cap.release()
    return frames

import supervision as sv
from typing import Any, Dict

def run_tracking(detections: List[Any]) -> Dict[str, List[Dict[int, Dict[str, List[float]]]]]:
    tracker = sv.ByteTrack()
    results = {"players": [], "referees": [], "ball": []}

    for det in detections:
        cls_names = det.names
        name2id = {v: k for k, v in cls_names.items()}
        sv_det = sv.Detections.from_ultralytics(det)

        for i, cls_id in enumerate(sv_det.class_id):
            if cls_names[cls_id] == "goalkeeper":
                sv_det.class_id[i] = name2id["player"]

        tracked = tracker.update_with_detections(sv_det)

        frame_result = {"players": {}, "referees": {}, "ball": {}}
        for t in tracked:
            bbox, cls_id, track_id = t[0].tolist(), t[3], t[4]
            if cls_id == name2id["player"]:
                frame_result["players"][track_id] = {"bbox": bbox}
            elif cls_id == name2id["referee"]:
                frame_result["referees"][track_id] = {"bbox": bbox}

        for t in sv_det:
            if t[3] == name2id["ball"]:
                frame_result["ball"][1] = {"bbox": t[0].tolist()}

        for key in results:
            results[key].append(frame_result[key])
    return results

# 图像预处理
upper_half = frame[0:int(h/2), :, :]  # 获取图像上半部分
pixels = upper_half.reshape(-1, 3)

# 聚类为两个颜色区域（假设为球衣+背景）
kmeans = KMeans(n_clusters=2).fit(pixels)

# 找出背景标签（角落像素最多的聚类中心）
corner_pixels = frame[0:10, 0:10, :].reshape(-1, 3)
corner_labels = kmeans.predict(corner_pixels)
background_label = np.bincount(corner_labels).argmax()

# 得到球员球衣颜色中心
player_cluster_color = kmeans.cluster_centers_[1 - background_label]

features = extract_color_features(frame, box)
tracker.update(box, features)

# 示例：两帧之间线性插值
ball_pos_frame1 = (x1, y1)
ball_pos_frame2 = (x2, y2)
interpolated = [(x1 + i*(x2-x1)/n, y1 + i*(y2-y1)/n) for i in range(1, n)]