冰球运动的AI科技感：用计算机视觉跟踪球员-51CTO.COM

译者 | 朱先忠

审校 | 重楼

引言

如今，我并没有像我孩提梦想的那样成为一名冰球运动员；但是，从我小时候起，冰球就成为我生活的一部分。最近，我有机会在秘鲁首都利马举行的第一届冰球锦标赛（3对3）中成为裁判的技术支持人员，并保存下来一些统计数据。当然，这项工作的完满完成还涉及到秘鲁直排冰球协会（APHL）的努力和一次由友谊联盟成功组织的友好访问。

为了增加人工智能技术的应用测试，我使用了PyTorch、计算机视觉技术和卷积神经网络（CNN）来构建一个模型，该模型可以跟踪球员和球队，并收集一些基本的表现统计数据。

本文旨在为设计和部署这一类模型提供一个快速指南。虽然这个模型还需要一些微调，但我希望它能帮助任何人了解应用于体育领域的计算机视觉的有趣世界。此外，我要感谢秘鲁直排冰球协会（APHL）允许我在这个项目中使用比赛的40秒视频样本（你可以在项目的GitHub存储库中找到该视频的输入样本）。

系统架构

在着手开发这个项目之前，我做了一些研究，以找到一个基础型框架，方便直接拿过来使用，从而避免“重新发明轮子”。我发现，在使用计算机视觉追踪球员方面，已经存在很多与足球有关的有趣工作（这并不奇怪，因为足球是世界上最受欢迎的团队运动）。然而，我没有找到很多的与冰球有关的资源。还好，我发现Roboflow提供了一些有趣的预训练模型和数据集用于训练开发者自己的模型，但他们使用了托管模型技术，从而导致会出现一些延迟问题，我将在文章稍后作进一步的解释。

最后，我按照本教程中解释的基本原理和跟踪方法，利用足球材料来读取视频帧并获取单个跟踪ID（如果你有兴趣更好地了解一些基本的计算机视觉技术，我建议你至少观看本教程的前一个半小时，教程地址：https://youtu.be/neBZ6huolkg?feature=shared）。

在完成跟踪ID标注后，我构建出自己的开发线路。在阅读本文时，我们将看到该项目如何从一个简单的对象检测任务演变为一个完整的能够检测球员、团队并提供一些基本性能指标的模型（文章中嵌入的所有从01到08编号的样本片段，都是作者自己创建的）。

AI冰球应用程序模型架构

跟踪机制

跟踪机制是本文构建模型的支柱部分。它确保视频中每个检测到的对象都被识别并分配一个唯一的标识符，在每一帧中都保持这种标识。具体说来，跟踪机制的主要组成部分包括以下组件：

YOLO算法：这是一种强大的实时目标检测算法，最初于2015年在《You Only Look:Unified，real time object detection》一文中介绍。它在检测大约80个预训练种类的速度和多功能性方面脱颖而出（值得注意的是，它也可以在自定义数据集上训练以检测特定对象）。对于我们的使用场景来说，我们将依赖YOLOv8x，这是Ultralytics.com公司基于之前的YOLO版本构建的计算机视觉模型。你可以在链接https://github.com/ultralytics/ultralytics处下载。
ByteTrack跟踪器：要理解ByteTrack，我们必须先了解一下MOT（多对象跟踪），它涉及跟踪视频序列中多个对象随时间的移动，并将当前帧中检测到的对象与前一帧中的相应对象链接起来。为了实现这一点，我们将使用ByteTrack（在2021年的论文《ByteTrack：通过关联每个检测框进行多目标跟踪》中介绍）。为了实现ByteTrack跟踪器并为检测到的对象分配跟踪ID，我们将依赖基于Python的Supervision库（［译者注］这是一款出色的开源的基于Python的计算机视觉低代码工具，其设计初衷是为用户提供一个便捷且高效的接口，用以处理数据集并直观地展示检测结果）。
OpenCV库：这是Python中用于各种计算机视觉任务的知名的开源库。对于我们的使用场景来说，我们将依靠OpenCV库实现可视化和注释视频帧，并为每个检测到的对象添加边界框和文本。

为了构建我们的跟踪机制，我们将从以下两个步骤开始：

使用ByteTrack部署YOLO模型来检测对象（在我们的例子中是球员）并分配唯一的跟踪ID。
初始化字典，以便将对象轨迹存储在pickle（pkl）文件中。这是非常有用的，因为这可以避免每次运行代码时都执行逐帧视频对象检测过程，并节省大量时间。

我们需要安装以下Python包：

pip install ultralytics
pip install supervision
pip install opencv-python1.
2.
3.

接下来，我们将指定我们的库以及示例视频文件和pickle文件的路径（如果存在的话；如果不存在，代码将创建一个新的文件并将其保存在相同的路径中）：

#**********************************库*********************************#
from ultralytics import YOLO
import supervision as sv
import pickle
import os
import cv2

#输入：视频文件
video_path = 'D:/PYTHON/video_input.mp4'
#输出：视频文件
output_video_path = 'D:/PYTHON/output_video.mp4'
# PICKLE文件（如果存在的话；如果不存在，代码将创建一个新的文件并将其保存在相同的路径中）
pickle_path = 'D:/PYTHON/stubs/track_stubs.pkl'1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.

现在，让我们继续定义我们的跟踪机制（你可以在项目的GitHub存储库中找到此视频输入示例）：

#*********************************跟踪机制**************************#
class HockeyAnalyzer:
def __init__(self, model_path):
self.model = YOLO(model_path) 
self.tracker = sv.ByteTrack()

def detect_frames(self, frames):
batch_size = 20 
detections = [] 
for i in range(0, len(frames), batch_size):
detections_batch = self.model.predict(frames[i:i+batch_size], conf=0.1)
detections += detections_batch
return detections

#********从文件加载轨迹或检测对象-保存PICKLE文件************#

def get_object_tracks(self, frames, read_from_stub=False, stub_path=None):
if read_from_stub and stub_path is not None and os.path.exists(stub_path):
with open(stub_path, 'rb') as f:
tracks = pickle.load(f)
return tracks

detections = self.detect_frames(frames)

tracks = {"person": []}

for frame_num, detection in enumerate(detections):
cls_names = detection.names
cls_names_inv = {v: k for k, v in cls_names.items()}

# 跟踪机制
detection_supervision = sv.Detections.from_ultralytics(detection)
detection_with_tracks = self.tracker.update_with_detections(detection_supervision)
tracks["person"].append({})

for frame_detection in detection_with_tracks:
bbox = frame_detection[0].tolist()
cls_id = frame_detection[3]
track_id = frame_detection[4]

if cls_id == cls_names_inv.get('person', None):
tracks["person"][frame_num][track_id] = {"bbox": bbox}

for frame_detection in detection_supervision:
bbox = frame_detection[0].tolist()
cls_id = frame_detection[3]

if stub_path is not None:
with open(stub_path, 'wb') as f:
pickle.dump(tracks, f)

return tracks

#***********************边界框与跟踪ID**************************#

def draw_annotations(self, video_frames, tracks):
output_video_frames = []
for frame_num, frame in enumerate(video_frames):
frame = frame.copy() 
player_dict = tracks["person"][frame_num]

# 绘制球员
for track_id, player in player_dict.items():
color = player.get("team_color", (0, 0, 255))  
bbox = player["bbox"]
x1, y1, x2, y2 = map(int, bbox)         
# 边界框
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
# Track_id 
cv2.putText(frame, str(track_id), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

output_video_frames.append(frame)

return output_video_frames1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.

上面的方法中，首先初始化YOLO模型和ByteTrack跟踪器。接下来，每帧以20个为一批次进行处理，使用YOLO模型检测和收集每个批次中的对象。如果pickle文件在其路径中可用，则它会从文件中预计算轨迹。如果pickle文件不可用（你是第一次运行代码或删除了之前的pickle文件），get_object_tracks会将每个检测转换为ByteTrack所需的格式，用这些检测更新跟踪器，并将跟踪信息存储在指定路径中的新pickle文件中。最后，对每一帧进行迭代，为每个检测到的对象绘制边界框和跟踪ID。

要执行跟踪器并保存带有边界框和跟踪ID的新输出视频，可以使用以下代码：

#*************** 执行跟踪机制并输出视频****************#

#读取视频帧
video_frames = []
cap = cv2.VideoCapture(video_path)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
video_frames.append(frame)
cap.release()

#********************* 用YOLO执行跟踪方法**********************#
tracker = HockeyAnalyzer('D:/PYTHON/yolov8x.pt')
tracks = tracker.get_object_tracks(video_frames, read_from_stub=True, stub_path=pickle_path)
annotated_frames = tracker.draw_annotations(video_frames, tracks)

#***********************保存视频文件************************************#
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
height, width, _ = annotated_frames[0].shape
out = cv2.VideoWriter(output_video_path, fourcc, 30, (width, height))

for frame in annotated_frames:
out.write(frame)
out.release()1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.

如果代码中的所有内容都正常工作，你应该会观察到一个类似于示例剪辑01中所示的视频输出。

示例剪辑01：基本跟踪机制（对象和跟踪ID）

【Tip-01】不要低估你的计算机的计算能力！第一次运行代码时，预计帧处理需要一些时间，这具体取决于你的计算机的计算能力。对我来说，只使用CPU设置需要45到50分钟（考虑CUDA作为选项）。YOLOv8x跟踪机制虽然强大，但需要大量的计算资源（有时，我的内存达到99%，但愿它不会崩溃！）。如果你在使用此版本的YOLO时遇到问题，Ultralytics的GitHub上提供了更轻量级的模型，以平衡准确性和计算能力。

溜冰场

正如你从本文一开始所看到的，开发中我们面临着不少挑战。首先，正如预期的那样，上面创建的模型拾取了所有运动物体——球员、裁判，甚至是溜冰场外的人。其次，那些红色的边界框可能会使跟踪球员有点不清楚，也不利于演示。在本节中，我们将专注于将检测范围缩小到溜冰场内的物体。此外，我们将把这些边界框替换为底部的椭圆，以确保更清晰的可见性。

首先，让我们先从使用方框切换到使用椭圆。为了实现这一点，我们只需在现有代码中的标签和边界框上方添加一个新方法：

#************用椭圆代替边界框来跟踪球员的设计**************#

def draw_ellipse(self, frame, bbox, color, track_id=None, team=None):
y2 = int(bbox[3])
x_center = (int(bbox[0]) + int(bbox[2])) // 2
width = int(bbox[2]) - int(bbox[0])
color = (255, 0, 0)
text_color = (255, 255, 255)

cv2.ellipse(
frame,
center=(x_center, y2),
axes=(int(width) // 2, int(0.35 * width)),
angle=0.0,
startAngle=-45,
endAngle=235,
color=color,
thickness=2,
lineType=cv2.LINE_4
)

if track_id is not None:
rectangle_width = 40
rectangle_height = 20
x1_rect = x_center - rectangle_width // 2
x2_rect = x_center + rectangle_width // 2
y1_rect = (y2 - rectangle_height // 2) + 15
y2_rect = (y2 + rectangle_height // 2) + 15

cv2.rectangle(frame,
(int(x1_rect), int(y1_rect)),
(int(x2_rect), int(y2_rect)),
color,
cv2.FILLED)

x1_text = x1_rect + 12
if track_id > 99:
x1_text -= 10
font_scale = 0.4
cv2.putText(
frame,
f"{track_id}",
(int(x1_text), int(y1_rect + 15)),
cv2.FONT_HERSHEY_SIMPLEX,
font_scale,
text_color,
thickness=2
)

return frame1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.

然后，我们还需要通过调用椭圆方法替换边界框和ID来更新注释步骤：

#***********************边界框和跟踪ID**************************#

def draw_annotations(self, video_frames, tracks):
output_video_frames = []
for frame_num, frame in enumerate(video_frames):
frame = frame.copy() 
player_dict = tracks["person"][frame_num]

# 绘制球员
for track_id, player in player_dict.items():
bbox = player["bbox"]

#绘制椭圆和跟踪ID
self.draw_ellipse(frame, bbox, (0, 255, 0), track_id)

x1, y1, x2, y2 = map(int, bbox)

output_video_frames.append(frame)

return output_video_frames1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.

通过上面这些更改，你的输出视频应该看起来更整洁一些，如示例剪辑02所示。

示例剪辑02：用椭圆替换边界框

现在，为了处理溜冰场的边界，我们需要对计算机视觉中的分辨率有一些基本的了解。在我们的示例中，我们使用的是720p（1280x720像素）格式；这意味着，我们处理的每一帧或图像的尺寸为1280像素（宽）乘720像素（高）。

使用720p（1280x720像素）格式意味着什么？这意味着，图像由水平1280像素和垂直720像素组成。此格式的坐标从图像左上角的（0,0）开始，x坐标随着向右移动而增加，y坐标随着向下移动而增加。这些坐标用于标记图像中的特定区域，例如将（x1，y1）用于框的左上角，将（x2，y2）用于框右下角。了解这一点将有助于我们测量距离和速度，并决定我们想在视频中集中分析的位置。

也就是说，我们将使用以下代码开始用绿线标记帧边界：

#********************* 帧的边界定义***********************
import cv2

video_path = 'D:/PYTHON/video_input.mp4'
cap = cv2.VideoCapture(video_path)

#**************读取、定义和绘制帧的角****************
ret, frame = cap.read()

bottom_left = (0, 720)
bottom_right = (1280, 720)
upper_left = (0, 0)
upper_right = (1280, 0)

cv2.line(frame, bottom_left, bottom_right, (0, 255, 0), 2)
cv2.line(frame, bottom_left, upper_left, (0, 255, 0), 2)
cv2.line(frame, bottom_right, upper_right, (0, 255, 0), 2)
cv2.line(frame, upper_left, upper_right, (0, 255, 0), 2)

#*******************保存带有标记角的帧*********************
output_image_path = 'rink_area_marked_VALIDATION.png'
cv2.imwrite(output_image_path, frame)
print("Rink area saved:", output_image_path)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.

结果应该是一个绿色矩形，如示例片段03中的（a）所示。但是，为了只跟踪溜冰场内的运动物体，我们需要一个更类似于（b）中的划界。

图03：溜冰场的边界定义

正确地得到（b）中的划界就像一个反复试验的过程。在这个过程中，你需要测试不同的坐标，直到找到最适合你的模型的边界。起初，我的目标是精确地匹配溜冰场的边界。然而，发现跟踪系统在溜冰场的边缘附近判别比较困难。为了提高准确性，我稍微扩大了一下边界，以确保捕捉到溜冰场内的所有跟踪对象，同时排除场外的跟踪对象。最后，如（b）中显示的结果是我所能得到的最好的结果（你仍然可以在更好的情况下工作），这个边界由下面这些关键坐标定义：

左下角：（-450，710）
右下角：（2030，710）
左上角：（352，61）
右上角：（948，61）

最后，我们将定义另外两个区域：白队和黄队的进攻区（每支球队的目标都是得分）。这将使我们能够收集对手区域内每支球队的一些基本位置统计数据和压力指标。

图04：进攻区

#**************黄队进攻区****************
Bottom Left Corner: (-450, 710)
Bottom Right Corner: (2030, 710)
Upper Left Corner: (200, 150)
Upper Right Corner: (1160, 150)

#**************白队进攻区****************
Bottom Left Corner: (180, 150)
Bottom Right Corner: (1100, 150)
Upper Left Corner: (352, 61)
Upper Right Corner: (900, 61)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.

我们将暂时搁置这些坐标，并在下一步解释我们将如何对每个团队进行分类。然后，我们将把它们整合到我们最初的跟踪方法中。

使用深度学习进行团队预测

自Warren McCulloch和Walter Pitts于1943年发表论文《神经活动中内在思想的逻辑演算》（https://www.cs.cmu.edu/~./epxing/Class/10715/reading/McCulloch.and.Pitts.pdf）以来，已经过去了80多年，该论文为早期神经网络研究奠定了坚实的基础。后来，在1957年，一个简化神经元的数学模型（接收输入，对这些输入应用权重，求和并输出二进制结果）启发了Frank Rosenblatt构建Mark I（https://news.cornell.edu/stories/2019/09/professors-perceptron-paved-way-ai-60-years-too-soon）。这是第一个旨在演示感知器概念的硬件实现，感知器是一种能够从数据中学习以进行二进制分类的神经网络模型。从那时起，让计算机像我们一样思考的探索就没有放缓。如果这是你第一次深入学习神经网络，或者你想更新和加强你的知识，那么我建议你阅读Shreya Rao的系列文章（https://medium.com/@shreya.rao/list/deep-learning-illustrated-ae6c27de1640），这可以作为深度学习的一个很好的起点。此外，你可以访问我链接处https://medium.com/@raul.vizcarrach/list/neural-networks-098e9b594f19收集的故事集（来自于不同的投稿者），你可能会发现这些故事很有用。

那么，为什么选择卷积神经网络（CNN）呢？老实说，这不是我的第一选择。最初，我尝试使用LandingAI构建一个模型，LandingAI是一个用户友好的云部署平台，通过API连接Python。然而，这一方案出现了延迟问题（在线处理超过1000帧）。尽管Roboflow中的预训练模型具有高质量的数据集和预训练模型，但它们也出现了类似的延迟问题。意识到需要在本地运行它，我尝试了一种基于MSE的方法来对球衣颜色进行分类，以便进行球队和裁判检测。虽然这听起来像是最终的解决方案，但它的准确性很低。经过几天的反复试验，我最终改用CNN。在不同的深度学习方法中，CNN非常适合对象检测，不像LSTM或RNN更适合语言转录或翻译等顺序数据。

在深入分析代码之前，让我们先了解一下CNN架构的一些基本概念：

学习样本数据集：数据集分为三类：裁判、Team_Away（白色球衣球员）和Team_Home（黄色球衣球员）。每个类别的样本被分为两组：训练数据和验证数据。CNN将在每个训练轮次中使用训练数据来“学习”多层模式。验证数据将在每次迭代结束时用于评估模型的性能，并衡量其对新数据的泛化程度。创建样本数据集并不难；我花了大约30到40分钟从视频中裁剪每个类别的示例图像，并将其组织到子目录中。我设法创建了一个大约90张图片的示例数据集，你可以在项目的GitHub存储库中找到。
模型是如何学习的？：输入数据在神经网络的每一层中移动，神经网络可以有一层或多层连接在一起进行预测。每一层都使用一个激活函数来处理数据，以进行预测或对数据进行更改。这些层之间的每个连接都有一个权重，它决定了一个层的输出对下一个层有多大的影响。目标是找到这些权重的正确组合，以尽量减少预测结果时的错误。通过称为反向传播和损失函数的过程，该模型调整这些权重以减少误差并提高准确性。这个过程在所谓的训练轮次（前向传递+反向传播）中重复，随着模型从错误中学习，它在每个周期的预测能力越来越好。
激活函数：如前所述，激活函数在模型的学习过程中起着重要作用。我选择了ReLU（校正线性单元）算法，因为它以计算效率高和缓解所谓的消失梯度问题（多层网络可能会有效地停止学习）而闻名。虽然ReLU工作良好，但其他激活函数如sigmoid、tanh或swish等也各有其用途，具体取决于网络的复杂程度。
训练轮次（epoch）：设定正确的训练轮次需要实验。你应该考虑数据集的复杂性、CNN模型的架构和计算资源等因素。在大多数情况下，最好在每次迭代中监控模型的性能，并在改进变得最小时停止训练，以防止过拟合。考虑到我的小训练数据集，我决定以10个训练轮次作为基线。然而，在其他情况下，可能需要根据指标性能和验证结果进行调整。
Adam（自适应矩估计）：最终，目标是减少预测输出和真实输出之间的误差。如前所述，反向传播在这里起着关键作用，它通过调整和更新神经网络权重来随着时间的推移改进预测。虽然反向传播基于损失函数的梯度处理权重更新，但Adam算法通过动态调整学习率来逐步最小化误差或损失函数，从而增强了这一过程。换句话说，它可以微调模型的学习速度。

也就是说，为了运行我们的CNN模型，我们需要以下Python包：

pip install torch torchvision 
pip install matplotlib 
pip install scikit-learn1.
2.
3.

【Tip-02】确保PyTorch安装正确。我所有的工具都是在Anaconda环境中设置的，当我安装PyTorch时，起初它似乎设置得很正确。然而，在运行一些库时却出现了一些问题。起初，我以为这是代码原因，但经过几次修改都没有成功，最后我不得不重新安装Anaconda并在干净的环境中重新安装了PyTorch。最终，问题就解决了！

接下来，我们将指定我们的库和样本数据集的路径：

# ************卷积神经网络三类检测**************************
# 裁判
# 白队(Team_away)
# 黄队 (Team_home)

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import matplotlib.pyplot as plt

#培训和验证数据集
#从项目的GitHub存储库下载teams_sample_dataset文件
data_dir = 'D:/PYTHON/teams_sample_dataset'1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.

首先，我们将确保每张图片大小相等（调整为150x150像素），然后将其转换为代码可以理解的格式（在PyTorch中，输入数据通常表示为Tensor对象）。最后，我们将调整颜色，使模型更容易使用（归一化），并设置一个加载图像的过程。这些步骤共同帮助准备图片并对其进行组织，以便模型能够有效地开始从中学习，避免数据格式引起的偏差。

#******************************数据转换***********************************
transform = transforms.Compose([
transforms.Resize((150, 150)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

#加载数据集
train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transform=transform)
val_dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'), transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.

接下来，我们将定义CNN的架构：

#********************************CNN模型架构**************************************
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.fc1 = nn.Linear(128 * 18 * 18, 512)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, 3)  #三个类（裁判，白队，黄队）

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
x = x.view(-1, 128 * 18 * 18)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)  
return x1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.

我们的CNN模型有三层（conv1、conv2、conv3）。数据从卷积层（conv）开始，在那里应用了激活函数（ReLU）。此函数使网络能够学习数据中的复杂模型和关系。随后，池化层被激活。什么是最大池化？这是一种在保留重要特征的同时减小图像大小的技术，有助于高效训练和优化内存资源。这个过程在conv1到conv3之间重复。最后，数据通过完全连接的层（fc1、fc2）进行最终分类（或决策）。

下一步，我们初始化模型，将分类交叉熵配置为损失函数（通常用于分类任务），并指定Adam作为我们的优化器。如前所述，我们将在10个训练轮次的完整周期内执行我们的模型。

#********************************CNN训练**********************************************

# 模型损失函数－优化器
model = CNNModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

#*********************************训练*************************************************
num_epochs = 10
train_losses, val_losses = [], []

for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
labels = labels.type(torch.LongTensor)  
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()

train_losses.append(running_loss / len(train_loader))

model.eval()
val_loss = 0.0
all_labels = []
all_preds = []
with torch.no_grad():
for inputs, labels in val_loader:
outputs = model(inputs)
labels = labels.type(torch.LongTensor)  
loss = criterion(outputs, labels)
val_loss += loss.item()
_, preds = torch.max(outputs, 1)  
all_labels.extend(labels.tolist())
all_preds.extend(preds.tolist())1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.

为了跟踪性能，我们将添加一些代码来跟踪训练进度，打印验证指标并绘制出来。最后，我们将模型保存为hockey_team_classifier.cth，并保存在你选择的指定路径中。

#********************************指标性能************************************

val_losses.append(val_loss / len(val_loader))
val_accuracy = accuracy_score(all_labels, all_preds)
val_precision = precision_score(all_labels, all_preds, average='macro', zero_division=1)
val_recall = recall_score(all_labels, all_preds, average='macro', zero_division=1)
val_f1 = f1_score(all_labels, all_preds, average='macro', zero_division=1)

print(f"Epoch [{epoch + 1}/{num_epochs}], "
f"Loss: {train_losses[-1]:.4f}, "
f"Val Loss: {val_losses[-1]:.4f}, "
f"Val Acc: {val_accuracy:.2%}, "
f"Val Precision: {val_precision:.4f}, "
f"Val Recall: {val_recall:.4f}, "
f"Val F1 Score: {val_f1:.4f}")

#*******************************显示指标 &性能**********************************
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Validation Loss')
plt.legend()
plt.show()

# 保存GH_CV_track_teams代码的模型
torch.save(model.state_dict(), 'D:/PYTHON/hockey_team_classifier.pth')1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.

此外，在运行完上述所有步骤（你可以在项目的GitHub存储库中找到完整的代码）后，除了你的“pth”文件外，你还应该看到以下输出（指标可能略有不同）：

图05：CNN模型性能指标

#**************CNN PERFORMANCE ACROSS TRAINING EPOCHS************************

Epoch [1/10], Loss: 1.5346, Val Loss: 1.2339, Val Acc: 47.37%, Val Precision: 0.7172, Val Recall: 0.5641, Val F1 Score: 0.4167
Epoch [2/10], Loss: 1.1473, Val Loss: 1.1664, Val Acc: 55.26%, Val Precision: 0.6965, Val Recall: 0.6296, Val F1 Score: 0.4600
Epoch [3/10], Loss: 1.0139, Val Loss: 0.9512, Val Acc: 57.89%, Val Precision: 0.6054, Val Recall: 0.6054, Val F1 Score: 0.5909
Epoch [4/10], Loss: 0.8937, Val Loss: 0.8242, Val Acc: 60.53%, Val Precision: 0.7222, Val Recall: 0.5645, Val F1 Score: 0.5538
Epoch [5/10], Loss: 0.7936, Val Loss: 0.7177, Val Acc: 63.16%, Val Precision: 0.6667, Val Recall: 0.6309, Val F1 Score: 0.6419
Epoch [6/10], Loss: 0.6871, Val Loss: 0.7782, Val Acc: 68.42%, Val Precision: 0.6936, Val Recall: 0.7128, Val F1 Score: 0.6781
Epoch [7/10], Loss: 0.6276, Val Loss: 0.5684, Val Acc: 78.95%, Val Precision: 0.8449, Val Recall: 0.7523, Val F1 Score: 0.7589
Epoch [8/10], Loss: 0.4198, Val Loss: 0.5613, Val Acc: 86.84%, Val Precision: 0.8736, Val Recall: 0.8958, Val F1 Score: 0.8653
Epoch [9/10], Loss: 0.3959, Val Loss: 0.3824, Val Acc: 92.11%, Val Precision: 0.9333, Val Recall: 0.9213, Val F1 Score: 0.9243
Epoch [10/10], Loss: 0.2509, Val Loss: 0.2651, Val Acc: 97.37%, Val Precision: 0.9762, Val Recall: 0.9792, Val F1 Score: 0.97691.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.

在完成10个迭代周期后，CNN模型的性能指标有所改善。最初，在第一个训练轮次中，该模型的训练损失为1.5346，验证准确率为47.37%。我们应该如何理解这种初始训练结果呢？

准确性是评估分类性能的最常见指标之一。在我们的例子中，它代表了正确预测的类别在总数中所占的比例。然而，仅靠高精度并不能保证整体模型性能；你仍然可能对特定类别的预测得到很差的预测结果（正如我在早期试验中所经历的那样）。关于训练损失函数，它衡量模型学习将输入数据映射到正确标签的有效性。由于我们使用的是分类函数，交叉熵损失量化了预测的类别概率和实际标签之间的差异。像1.5346这样的起始值表示预测类别和实际类别之间存在显著差异；理想情况下，随着训练的进行，该值应接近0。随着时间的推移，我们观察到训练损失显著下降，验证准确性提高。到最后一个训练轮次，训练和验证损失分别达到0.2509和0.2651的低点。

为了测试我们的CNN模型，我们可以选择一个球员图像样本并评估其预测能力。为了进行测试，你可以运行以下代码并使用项目GitHub存储库中的validation_dataset文件夹。

# *************用样本数据集测试CNN模型***************************

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image

# 用于验证的样本数据集
test_dir = 'D:/PYTHON/validation_dataset'

# 团队预测的CNN模型
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.fc1 = nn.Linear(128 * 18 * 18, 512)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, 3) 

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
x = x.view(-1, 128 * 18 * 18)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)  
return x

# 之前保存的CNN模型
model = CNNModel()
model.load_state_dict(torch.load('D:/PYTHON/hockey_team_classifier.pth'))
model.eval()

transform = transforms.Compose([
transforms.Resize((150, 150)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

#******************样本图像的迭代精度检验*****************************

class_names = ['team_referee', 'team_away', 'team_home']

def predict_image(image_path, model, transform):
# 加载数据集
image = Image.open(image_path)
image = transform(image).unsqueeze(0)  

# 进行预测
with torch.no_grad():
output = model(image)
_, predicted = torch.max(output, 1)  
team = class_names[predicted.item()]
return team

for image_name in os.listdir(test_dir):
image_path = os.path.join(test_dir, image_name)
if os.path.isfile(image_path):  
predicted_team = predict_image(image_path, model, transform)
print(f'Image {image_name}: The player belongs to {predicted_team}')1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.

输出应该看起来像下面这样：

# *************CNN MODEL TEST - OUTPUT ***********************************#

Image Away_image04.jpg: The player belongs to team_away
Image Away_image12.jpg: The player belongs to team_away
Image Away_image14.jpg: The player belongs to team_away
Image Home_image07.jpg: The player belongs to team_home
Image Home_image13.jpg: The player belongs to team_home
Image Home_image16.jpg: The player belongs to team_home
Image Referee_image04.jpg: The player belongs to team_referee
Image Referee_image09.jpg: The player belongs to team_referee
Image Referee_image10.jpg: The player belongs to team_referee
Image Referee_image11.jpg: The player belongs to team_referee1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.

正如你所看到的，该模型在识别球队和排除裁判作为团队球员方面表现出了很好的能力。

【提示03】我在CNN设计过程中学到的一点是，增加复杂性并不总是能提高性能。最初，我尝试了更深层次的模型（更多的卷积层）和基于颜色的增强，以提高球员的球衣识别率。然而，在我的小数据集中，我遇到了过拟合，而不是学习可泛化的特征（所有图像都被预测为白人球队球员或裁判）。像dropout和批归一化这样的正则化技术也很重要；它们有助于在训练过程中施加约束，确保模型能够很好地泛化到新数据。就结果而言，更少有时意味着更多。

组合应用

将所有上面这些技术放在一起应用时需要对前面描述的跟踪机制进行一些调整。下面描述更新代码的逐步分解过程。

首先，我们将设置所需的库和路径。请注意，现在指定了pickle文件和CNN模型的路径。这一次，如果在路径中找不到pickle文件，代码将抛出错误。如果需要，使用前面的代码生成pickle文件，并使用此更新版本执行视频分析：

import cv2
import numpy as np
from ultralytics import YOLO
import pickle
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image

# 模型输入
model_path = 'D:/PYTHON/yolov8x.pt'
video_path = 'D:/PYTHON/video_input.mp4'
output_path = 'D:/PYTHON/output_video.mp4'
tracks_path = 'D:/PYTHON/stubs/track_stubs.pkl'
classifier_path = 'D:/PYTHON/hockey_team_classifier.pth'1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.

接下来，我们将加载模型，指定溜冰场坐标，并像以前一样，以20个为一批启动检测每帧中对象的过程。请注意，目前，我们将只使用溜冰场边界来将分析重点放在溜冰场上。在本文的最后步骤中，当我们包含性能统计数据时，我们将使用进攻区域坐标。

#***************************加载模型和溜冰场坐标********************#
class_names = ['Referee', 'Tm_white', 'Tm_yellow']

class HockeyAnalyzer:
def __init__(self, model_path, classifier_path):
self.model = YOLO(model_path)
self.classifier = self.load_classifier(classifier_path)
self.transform = transforms.Compose([
transforms.Resize((150, 150)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
self.rink_coordinates = np.array([[-450, 710], [2030, 710], [948, 61], [352, 61]])
self.zone_white = [(180, 150), (1100, 150), (900, 61), (352, 61)]
self.zone_yellow = [(-450, 710), (2030, 710), (1160, 150), (200, 150)]

#********************检测每帧中的对象 **********************************#
def detect_frames(self, frames):
batch_size = 20 
detections = [] 
for i in range(0, len(frames), batch_size):
detections_batch = self.model.predict(frames[i:i+batch_size], conf=0.1)
detections += detections_batch
return detections1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.

接下来，我们将添加预测每个球员球队的过程：

#*********************** 加载CNN模型**********************************************#

def load_classifier(self, classifier_path):
model = CNNModel()
model.load_state_dict(torch.load(classifier_path, map_location=torch.device('cpu')))
model.eval()
return model

def predict_team(self, image):
with torch.no_grad():
output = self.classifier(image)
_, predicted = torch.max(output, 1)
predicted_index = predicted.item()
team = class_names[predicted_index]
return team1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.

下一步，我们将添加前面描述的方法，从边界框切换到椭圆：

#************ 使用椭圆来跟踪球员，而不是边界框*******************#
def draw_ellipse(self, frame, bbox, color, track_id=None, team=None):
y2 = int(bbox[3])
x_center = (int(bbox[0]) + int(bbox[2])) // 2
width = int(bbox[2]) - int(bbox[0])

if team == 'Referee':
color = (0, 255, 255)
text_color = (0, 0, 0)
else:
color = (255, 0, 0)
text_color = (255, 255, 255)

cv2.ellipse(
frame,
center=(x_center, y2),
axes=(int(width) // 2, int(0.35 * width)),
angle=0.0,
startAngle=-45,
endAngle=235,
color=color,
thickness=2,
lineType=cv2.LINE_4
)

if track_id is not None:
rectangle_width = 40
rectangle_height = 20
x1_rect = x_center - rectangle_width // 2
x2_rect = x_center + rectangle_width // 2
y1_rect = (y2 - rectangle_height // 2) + 15
y2_rect = (y2 + rectangle_height // 2) + 15

cv2.rectangle(frame,
(int(x1_rect), int(y1_rect)),
(int(x2_rect), int(y2_rect)),
color,
cv2.FILLED)

x1_text = x1_rect + 12
if track_id > 99:
x1_text -= 10
font_scale = 0.4
cv2.putText(
frame,
f"{track_id}",
(int(x1_text), int(y1_rect + 15)),
cv2.FONT_HERSHEY_SIMPLEX,
font_scale,
text_color,
thickness=2
)

return frame1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.

现在，是时候添加分析器了，包括读取pickle文件，将分析范围缩小到我们之前定义的溜冰场边界内，并调用CNN模型来识别每个球员的球队并添加标签。请注意，我们提供了一个函数，可以用不同的颜色标记裁判，并更改其椭圆的颜色。在代码的最后部分，将处理后的帧写入输出视频。

#*******************加载跟踪的数据 (pickle 文件)**********************************#

def analyze_video(self, video_path, output_path, tracks_path):
with open(tracks_path, 'rb') as f:
tracks = pickle.load(f)

cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print("Error: Could not open video.")
return

fps = cap.get(cv2.CAP_PROP_FPS)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))

frame_num = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break

#***********检查球员是否落入溜冰场区域**********************************#
mask = np.zeros(frame.shape[:2], dtype=np.uint8)
cv2.fillConvexPoly(mask, self.rink_coordinates, 1)
mask = mask.astype(bool)
# 绘制溜冰场边缘
#cv2.polylines(frame, [self.rink_coordinates], isClosed=True, color=(0, 255, 0), thickness=2)

# 从帧中获取跟踪信息
player_dict = tracks["person"][frame_num]
for track_id, player in player_dict.items():
bbox = player["bbox"]

# 检查球员是否落入溜冰场区域
x_center = int((bbox[0] + bbox[2]) / 2)
y_center = int((bbox[1] + bbox[3]) / 2)

if not mask[y_center, x_center]:
continue  

#**********************************球队预测********************************************#
x1, y1, x2, y2 = map(int, bbox)
cropped_image = frame[y1:y2, x1:x2]
cropped_pil_image = Image.fromarray(cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB))
transformed_image = self.transform(cropped_pil_image).unsqueeze(0)
team = self.predict_team(transformed_image)

#************ 使用椭圆追踪球员和标签******************************************#
self.draw_ellipse(frame, bbox, (0, 255, 0), track_id, team)

font_scale = 1  
text_offset = -20  

if team == 'Referee':
rectangle_width = 60
rectangle_height = 25
x1_rect = x1
x2_rect = x1 + rectangle_width
y1_rect = y1 - 30
y2_rect = y1 - 5
#针对裁判使用不同设置
cv2.rectangle(frame,
(int(x1_rect), int(y1_rect)),
(int(x2_rect), int(y2_rect)),
(0, 0, 0),  
cv2.FILLED)
text_color = (255, 255, 255)  
else:
if team == 'Tm_white':
text_color = (255, 215, 0)  #白队：蓝色标签
else:
text_color = (0, 255, 255)  #黄队：黄色标签

#绘制球队标签
cv2.putText(
frame,
team,
(int(x1), int(y1) + text_offset), 
cv2.FONT_HERSHEY_PLAIN,            
font_scale,
text_color,
thickness=2
)

# 写输出视频
out.write(frame)
frame_num += 1

cap.release()
out.release()1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.
80.
81.
82.
83.
84.
85.
86.
87.
88.
89.
90.
91.
92.
93.

最后，我们实现CNN的架构（在CNN设计过程中定义）并执行Hockey分析器：

#**********************CNN模型架构******************************#
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.fc1 = nn.Linear(128 * 18 * 18, 512)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, len(class_names))  

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
x = x.view(-1, 128 * 18 * 18)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)
return x

#*********执行HockeyAnalyzer/分类器并保存输出***********#
analyzer = HockeyAnalyzer(model_path, classifier_path)
analyzer.analyze_video(video_path, output_path, tracks_path)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.

运行所有步骤后，你的视频输出应该如下所示：

示例剪辑06：跟踪球员和团队

请注意，在最后一次更新中，物体检测仅在溜冰场内进行，球队和裁判都是不同的。虽然CNN模型仍需要微调，偶尔会对一些球员失去稳定性，但它在整个视频中仍然基本可靠和准确。

速度、距离和进攻压力

实验结果证明，本示例开发的模型方案在跟踪球队和球员信息方面展现出强大能力，这为衡量各种指标的应用开辟了令人兴奋的可能性，例如生成热图、分析速度和覆盖距离、跟踪区域入口或出口等动作，甚至还可以更深入地了解详细的球员指标。不妨让我们作一个测试，我们添加三个性能指标：每位球员的平均速度、每支球队的滑冰距离和进攻压力（以每支球队在对手区域内投入的距离值百分比来衡量）。我会把更详细的统计数据留给你！

最开始，我们将溜冰场的坐标从基于像素的测量值调整为基于米单位。这种调整使我们能够以米而不是像素读取数据。视频中看到的溜冰场的真实尺寸约为15mx30m（宽15米，高30米）。为了便于这种转换，我们引入了一种将像素坐标转换为米的方法。通过定义溜冰场的实际尺寸并使用其角的像素坐标（从左到右和从上到下），我们获得了转换因子。这些因素将支持我们估算以米为单位的距离和以米每秒为单位的速度的过程。（你可以探索和应用的另一种有趣的技术是透视变换）

#*********************加载模型和溜冰场坐标*****************#
class_names = ['Referee', 'Tm_white', 'Tm_yellow']

class HockeyAnalyzer:
def __init__(self, model_path, classifier_path):
*
*
*
*
*
*
self.pixel_to_meter_conversion() #<------ 添加这个工具方法

#***********将溜冰场的坐标从基于像素的测量值调整为基于米单位***************************#
def pixel_to_meter_conversion(self):
#溜冰场实际尺寸（单位：米）
rink_width_m = 15
rink_height_m = 30

#溜冰场尺寸的像素坐标
left_pixel, right_pixel = self.rink_coordinates[0][0], self.rink_coordinates[1][0]
top_pixel, bottom_pixel = self.rink_coordinates[2][1], self.rink_coordinates[0][1]

#转换因子
self.pixels_per_meter_x = (right_pixel - left_pixel) / rink_width_m
self.pixels_per_meter_y = (bottom_pixel - top_pixel) / rink_height_m

def convert_pixels_to_meters(self, distance_pixels):
#把像素转换成米
return distance_pixels / self.pixels_per_meter_x, distance_pixels / self1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.

我们现在准备为每个球员增加速度，以米每秒为单位。为此，我们需要进行三处修改。首先，在HockeyAnalyzer类中启动一个名为previous_positions的空字典，以帮助我们比较球员的当前和先前位置。同样，我们将创建一个team_stats结构来存储每个团队的统计数据，以便进一步可视化。

接下来，我们将添加一种速度方法来估计球员的速度（以每秒像素为单位），然后使用转换因子（前面解释过）将其转换为每秒米数。最后，根据analyze_video方法，我们将调用新的速度方法，并将速度添加到每个被跟踪的对象（球员和裁判）中。下面给出的是修改后的代码：

#*********************加载模型和溜冰场坐标*****************#
class_names = ['Referee', 'Tm_white', 'Tm_yellow']

class HockeyAnalyzer:
def __init__(self, model_path, classifier_path):
*
*
*
*
*
*
*
self.pixel_to_meter_conversion() 
self.previous_positions = {} #<------初始化空字典 
self.team_stats = {
'Tm_white': {'distance': 0, 'speed': [], 'count': 0, 'offensive_pressure': 0},
'Tm_yellow': {'distance': 0, 'speed': [], 'count': 0, 'offensive_pressure': 0}
} #<------ 初始化空字典

#**************** 速度：米每秒********************************#
def calculate_speed(self, track_id, x_center, y_center, fps):
current_position = (x_center, y_center)
if track_id in self.previous_positions:
prev_position = self.previous_positions[track_id]
distance_pixels = np.linalg.norm(np.array(current_position) - np.array(prev_position))
distance_meters_x, distance_meters_y = self.convert_pixels_to_meters(distance_pixels)
speed_meters_per_second = (distance_meters_x**2 + distance_meters_y**2)**0.5 * fps
else:
speed_meters_per_second = 0
self.previous_positions[track_id] = current_position
return speed_meters_per_second

#******************* 加载跟踪数据 (pickle 文件)**********************************#

def analyze_video(self, video_path, output_path, tracks_path):
with open(tracks_path, 'rb') as f:
tracks = pickle.load(f)

*
*
*
*
*
*
*
*
# 绘制球队标签
cv2.putText(
frame,
team,
(int(x1), int(y1) + text_offset), 
cv2.FONT_HERSHEY_PLAIN,            
font_scale,
text_color,
thickness=2
)

#**************添加下面几行代码--->:

speed = self.calculate_speed(track_id, x_center, y_center, fps)
#速度标签 
speed_font_scale = 0.8  
speed_y_position = int(y1) + 20
if speed_y_position > int(y1) - 5:
speed_y_position = int(y1) - 5

cv2.putText(
frame,
f"Speed: {speed:.2f} m/s",  
(int(x1), speed_y_position),  
cv2.FONT_HERSHEY_PLAIN,       
speed_font_scale,
text_color,
thickness=2
)

# 写输出视频
out.write(frame)
frame_num += 1

cap.release()
out.release()1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.
80.
81.
82.

如果你在添加上面这几行新代码时遇到问题，你可以随时访问该项目的GitHub存储库（https://github.com/rvizcarra15/IceHockey_ComputerVision_PyTorch），在那里你可以找到完整的参考代码。此时，你的视频输出结果应该像下图的样子（请注意，速度已添加到每个播放器的标签中）：

示例剪辑07：跟踪球员和速度

最后，让我们添加一个统计板，借助这个统计板我们可以跟踪每支球队每位球员的平均速度，以及其他指标，如覆盖距离和对手区域的进攻压力。

我们已经定义了进攻区，并将其整合到我们的代码中。现在，我们需要跟踪每个球员进入对手区域的频率。为了实现这一点，我们将使用光线投射算法实现一种方法。该算法检查球员的位置是否在白队或黄队的进攻区域内。它的工作原理是从球员到目标区域画一条假想线。如果线穿过一个边界，则球员在里面；如果它穿过更多边界（在我们的例子中，四个边界中的两个），则球员就在外面。然后，该代码扫描整个视频，以确定每个被跟踪对象的区域状态。

#************ 在目标区域中定位球员的位置***********************#

def is_inside_zone(self, position, zone):
x, y = position
n = len(zone)
inside = False
p1x, p1y = zone[0]
for i in range(n + 1):
p2x, p2y = zone[i % n]
if y > min(p1y, p2y):
if y <= max(p1y, p2y):
if x <= max(p1x, p2x):
if p1y != p2y:
xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
if p1x == p2x or x <= xinters:
inside = not inside
p1x, p1y = p2x, p2y
return inside1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.

现在，我们将通过添加一种方法来处理表现指标，该方法以表格形式显示每支球队的平均球员速度、总覆盖距离和进攻压力（在对手区域花费的时间百分比）。使用OpenCV，我们将把这些指标格式化为覆盖在视频上的表格，并引入动态更新机制，以便在游戏过程中保持实时统计数据。

#*******************************性能度量*********************************************#
def draw_stats(self, frame):
avg_speed_white = np.mean(self.team_stats['Tm_white']['speed']) if self.team_stats['Tm_white']['count'] > 0 else 0
avg_speed_yellow = np.mean(self.team_stats['Tm_yellow']['speed']) if self.team_stats['Tm_yellow']['count'] > 0 else 0
distance_white = self.team_stats['Tm_white']['distance']
distance_yellow = self.team_stats['Tm_yellow']['distance']

offensive_pressure_white = self.team_stats['Tm_white'].get('offensive_pressure', 0)
offensive_pressure_yellow = self.team_stats['Tm_yellow'].get('offensive_pressure', 0)

Pressure_ratio_W = offensive_pressure_white/distance_white   *100  if self.team_stats['Tm_white']['distance'] > 0 else 0
Pressure_ratio_Y = offensive_pressure_yellow/distance_yellow *100  if self.team_stats['Tm_yellow']['distance'] > 0 else 0

table = [
["", "Away_White", "Home_Yellow"],
["Average Speed\nPlayer", f"{avg_speed_white:.2f} m/s", f"{avg_speed_yellow:.2f} m/s"],
["Distance\nCovered", f"{distance_white:.2f} m", f"{distance_yellow:.2f} m"],
["Offensive\nPressure %", f"{Pressure_ratio_W:.2f} %", f"{Pressure_ratio_Y:.2f} %"],
]

text_color = (0, 0, 0)  
start_x, start_y = 10, 590  
row_height = 30     # 管理行间高度
column_width = 150  # 管理行间宽度
font_scale = 1  

def put_multiline_text(frame, text, position, font, font_scale, color, thickness, line_type, line_spacing=1.0):
y0, dy = position[1], int(font_scale * 20 * line_spacing)  # 在此处调整行距
for i, line in enumerate(text.split('\n')):
y = y0 + i * dy
cv2.putText(frame, line, (position[0], y), font, font_scale, color, thickness, line_type)

for i, row in enumerate(table):
for j, text in enumerate(row):
if i in [1,2, 3]:  
put_multiline_text(
frame,
text,
(start_x + j * column_width, start_y + i * row_height),
cv2.FONT_HERSHEY_PLAIN,
font_scale,
text_color,
1,
cv2.LINE_AA,
line_spacing= 0.8 
)
else:
cv2.putText(
frame,
text,
(start_x + j * column_width, start_y + i * row_height),
cv2.FONT_HERSHEY_PLAIN,
font_scale,
text_color,
1,
cv2.LINE_AA,
)       

#****************** 跟踪和更新游戏统计数据****************************************#

def update_team_stats(self, team, speed, distance, position):
if team in self.team_stats:
self.team_stats[team]['speed'].append(speed)
self.team_stats[team]['distance'] += distance
self.team_stats[team]['count'] += 1

if team == 'Tm_white':
if self.is_inside_zone(position, self.zone_white):
self.team_stats[team]['offensive_pressure'] += distance
elif team == 'Tm_yellow':
if self.is_inside_zone(position, self.zone_yellow):
self.team_stats[team]['offensive_pressure'] += distance
为了在视频中显示统计数据，我们必须调用analyze_video方法，所以一定要在定义速度标签之后，在处理输出视频之前添加下面这些额外的代码行：
*
*
*
*
*
*
*
#速度标签 
speed_font_scale = 0.8  
speed_y_position = int(y1) + 20
if speed_y_position > int(y1) - 5:
speed_y_position = int(y1) - 5

cv2.putText(
frame,
f"Speed: {speed:.2f} m/s",  
(int(x1), speed_y_position),  
cv2.FONT_HERSHEY_PLAIN,       
speed_font_scale,
text_color,
thickness=2
)
#**************添加下面几行代码--->:

distance = speed / fps
position = (x_center, y_center)
self.update_team_stats(team, speed, distance, position)

# 写输出视频
out.write(frame)
frame_num += 11.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.
80.
81.
82.
83.
84.
85.
86.
87.
88.
89.
90.
91.
92.
93.
94.
95.
96.
97.
98.
99.
100.
101.
102.
103.
104.

每个球员覆盖的距离（以米为单位）是通过将他们的速度（以米每秒为单位）除以帧率（每秒帧数）来计算的。这种计算使我们能够估计每个球员在视频中的每一帧变化之间移动了多远。如果一切正常，你的最终视频输出应该如下：

示例剪辑08：示例程序最终输出结果

待考虑因素和未来的工作

本文示例中实现的模型是使用计算机视觉跟踪冰球比赛（或任何团队运动）中有关球员的基本设置。但是，可以进行很多微调来改进它并添加新功能。以下是我为下一个2.0版本所做的一些想法，你也可以考虑：

跟踪冰球的挑战：考虑到冰球与足球或篮球相比的大小，跟踪冰球具有挑战性，这取决于你的相机朝向哪个方向和分辨率。但如果你实现了这一点，那么跟踪将表现出非常有趣的可能性，比如控球时间指标、进球机会或投篮数据。这也适用于个人表演；在冰球中，球员的变化比其他团队运动要频繁得多，因此跟踪每个球员在一个时期的表现是一个挑战。
计算资源，哦，为什么要计算呢？我在CPU上运行了所有代码，但遇到了问题（有时会导致蓝屏），由于在设计过程中内存不足（考虑使用CUDA设置）。我们的示例视频长约40秒，最初大小为5MB，但在运行模型后，输出增加到34MB。想象一下一个20分钟时间长的游戏的情况吧。因此，在扩展时，你应该考虑计算资源和存储。
不要低估MLOps：为了快速部署和扩展，我们需要高效、支持频繁执行且可靠的机器学习管道。这涉及到考虑采用持续集成部署训练方法。我们的用例是为特定场景构建的，但如果条件发生变化，比如相机方向或球衣颜色，该怎么办？为了扩大规模，我们必须采取CI/CD/CT的思维方式。

最后，我希望你觉得本文提供的这个计算机视觉演示项目非常有趣，你可以在GitHub存储库（https://github.com/rvizcarra15/IceHockey_ComputerVision_PyTorch）中访问完整的代码。如果你想支持该地区直排冰球和冰球的发展，请关注APHL（我们总是需要你想为年轻球员捐赠的二手设备，并致力于建设我们的第一个官方冰球场），在全球范围内，请关注并支持友谊联盟（https://friendshipleague.org/）。

原文标题：Spicing up Ice Hockey with AI: Player Tracking with Computer Vision，作者：Raul Vizcarra Chirinos

链接：https://towardsdatascience.com/spicing-up-ice-hockey-with-ai-player-tracking-with-computer-vision-ce9ceec9122a

想了解更多AIGC的内容，请访问：

51CTO AI.x社区

https://www.51cto.com/aigc/