冰球運動的AI科技感：用計算機視覺跟蹤球員 原創(chuàng)
本文將介紹使用PyTorch、計算機視覺庫OpenCV和卷積神經(jīng)網(wǎng)絡（CNN）開發(fā)一個跟蹤冰球球員、球隊及其基本表現(xiàn)統(tǒng)計數(shù)據(jù)的YOLOv8x模型的全過程。
引言
如今，我并沒有像我孩提夢想的那樣成為一名冰球運動員；但是，從我小時候起，冰球就成為我生活的一部分。最近，我有機會在秘魯首都利馬舉行的第一屆冰球錦標賽（3對3）中成為裁判的技術支持人員，并保存下來一些統(tǒng)計數(shù)據(jù)。當然，這項工作的完滿完成還涉及到秘魯直排冰球協(xié)會（APHL）的努力和一次由友誼聯(lián)盟成功組織的友好訪問。
為了增加人工智能技術的應用測試，我使用了PyTorch、計算機視覺技術和卷積神經(jīng)網(wǎng)絡（CNN）來構建一個模型，該模型可以跟蹤球員和球隊，并收集一些基本的表現(xiàn)統(tǒng)計數(shù)據(jù)。
本文旨在為設計和部署這一類模型提供一個快速指南。雖然這個模型還需要一些微調(diào)，但我希望它能幫助任何人了解應用于體育領域的計算機視覺的有趣世界。此外，我要感謝秘魯直排冰球協(xié)會（APHL：??https://www.instagram.com/aphl.pe/?igsh=MThvZWxhNThwdXpibA%3D%3D??）允許我在這個項目中使用比賽的40秒視頻樣本（你可以在項目的GitHub存儲庫https://github.com/rvizcarra15/IceHockey_ComputerVision_PyTorch中找到該視頻的輸入樣本）。
系統(tǒng)架構
在著手開發(fā)這個項目之前，我做了一些研究，以找到一個基礎型框架，方便直接拿過來使用，從而避免“重新發(fā)明輪子”。我發(fā)現(xiàn)，在使用計算機視覺追蹤球員方面，已經(jīng)存在很多與足球有關的有趣工作（這并不奇怪，因為足球是世界上最受歡迎的團隊運動）。然而，我沒有找到很多的與冰球有關的資源。還好，我發(fā)現(xiàn)Roboflow（https://universe.roboflow.com/search?q=hockey）提供了一些有趣的預訓練模型和數(shù)據(jù)集用于訓練開發(fā)者自己的模型，但他們使用了托管模型技術，從而導致會出現(xiàn)一些延遲問題，我將在文章稍后作進一步的解釋。
最后，我按照本教程中解釋的基本原理和跟蹤方法，利用足球材料來讀取視頻幀并獲取單個跟蹤ID（如果你有興趣更好地了解一些基本的計算機視覺技術，我建議你至少觀看本教程的前一個半小時，教程地址：https://youtu.be/neBZ6huolkg?feature=shared）。
在完成跟蹤ID標注后，我構建出自己的開發(fā)線路。在閱讀本文時，我們將看到該項目如何從一個簡單的對象檢測任務演變?yōu)橐粋€完整的能夠檢測球員、團隊并提供一些基本性能指標的模型（文章中嵌入的所有從01到08編號的樣本片段，都是作者自己創(chuàng)建的）。
AI冰球應用程序模型架構
跟蹤機制
跟蹤機制是本文構建模型的支柱部分。它確保視頻中每個檢測到的對象都被識別并分配一個唯一的標識符，在每一幀中都保持這種標識。具體說來，跟蹤機制的主要組成部分包括以下組件：
- YOLO算法：這是一種強大的實時目標檢測算法，最初于2015年在《You Only Look:Unified，real time object detection》（https://arxiv.org/abs/1506.02640）一文中介紹。它在檢測大約80個預訓練種類的速度和多功能性方面脫穎而出（值得注意的是，它也可以在自定義數(shù)據(jù)集上訓練以檢測特定對象）。對于我們的使用場景來說，我們將依賴YOLOv8x，這是Ultralytics.com公司基于之前的YOLO版本構建的計算機視覺模型。你可以在鏈接https://github.com/ultralytics/ultralytics處下載。
- ByteTrack跟蹤器：要理解ByteTrack，我們必須先了解一下MOT（多對象跟蹤），它涉及跟蹤視頻序列中多個對象隨時間的移動，并將當前幀中檢測到的對象與前一幀中的相應對象鏈接起來。為了實現(xiàn)這一點，我們將使用ByteTrack（在2021年的論文《ByteTrack：通過關聯(lián)每個檢測框進行多目標跟蹤》（https://arxiv.org/abs/2110.06864）中介紹）。為了實現(xiàn)ByteTrack跟蹤器并為檢測到的對象分配跟蹤ID，我們將依賴基于Python的Supervision庫（［譯者注］這是一款出色的開源的基于Python的計算機視覺低代碼工具，其設計初衷是為用戶提供一個便捷且高效的接口，用以處理數(shù)據(jù)集并直觀地展示檢測結果）。
- OpenCV庫：這是Python中用于各種計算機視覺任務的知名的開源庫。對于我們的使用場景來說，我們將依靠OpenCV庫（https://opencv.org/）實現(xiàn)可視化和注釋視頻幀，并為每個檢測到的對象添加邊界框和文本。
為了構建我們的跟蹤機制，我們將從以下兩個步驟開始：
- 使用ByteTrack部署YOLO模型來檢測對象（在我們的例子中是球員）并分配唯一的跟蹤ID。
- 初始化字典，以便將對象軌跡存儲在pickle（pkl）文件中。這是非常有用的，因為這可以避免每次運行代碼時都執(zhí)行逐幀視頻對象檢測過程，并節(jié)省大量時間。
我們需要安裝以下Python包：
pip install ultralytics
pip install supervision
pip install opencv-python接下來，我們將指定我們的庫以及示例視頻文件和pickle文件的路徑（如果存在的話；如果不存在，代碼將創(chuàng)建一個新的文件并將其保存在相同的路徑中）：
#**********************************庫*********************************#
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文件（如果存在的話；如果不存在，代碼將創(chuàng)建一個新的文件并將其保存在相同的路徑中）
pickle_path = 'D:/PYTHON/stubs/track_stubs.pkl'現(xiàn)在，讓我們繼續(xù)定義我們的跟蹤機制（你可以在項目的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_frames上面的方法中，首先初始化YOLO模型和ByteTrack跟蹤器。接下來，每幀以20個為一批次進行處理，使用YOLO模型檢測和收集每個批次中的對象。如果pickle文件在其路徑中可用，則它會從文件中預計算軌跡。如果pickle文件不可用（你是第一次運行代碼或刪除了之前的pickle文件），get_object_tracks會將每個檢測轉(zhuǎn)換為ByteTrack所需的格式，用這些檢測更新跟蹤器，并將跟蹤信息存儲在指定路徑中的新pickle文件中。最后，對每一幀進行迭代，為每個檢測到的對象繪制邊界框和跟蹤ID。
要執(zhí)行跟蹤器并保存帶有邊界框和跟蹤ID的新輸出視頻，可以使用以下代碼：
#*************** 執(zhí)行跟蹤機制并輸出視頻****************#
#讀取視頻幀
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執(zhí)行跟蹤方法**********************#
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()如果代碼中的所有內(nèi)容都正常工作，你應該會觀察到一個類似于示例剪輯01中所示的視頻輸出。
示例剪輯01：基本跟蹤機制（對象和跟蹤ID）
【Tip-01】不要低估你的計算機的計算能力！第一次運行代碼時，預計幀處理需要一些時間，這具體取決于你的計算機的計算能力。對我來說，只使用CPU設置需要45到50分鐘（考慮CUDA作為選項）。YOLOv8x跟蹤機制雖然強大，但需要大量的計算資源（有時，我的內(nèi)存達到99%，但愿它不會崩潰！）。如果你在使用此版本的YOLO時遇到問題，Ultralytics的GitHub上提供了更輕量級的模型，以平衡準確性和計算能力。
溜冰場
正如你從本文一開始所看到的，開發(fā)中我們面臨著不少挑戰(zhàn)。首先，正如預期的那樣，上面創(chuàng)建的模型拾取了所有運動物體——球員、裁判，甚至是溜冰場外的人。其次，那些紅色的邊界框可能會使跟蹤球員有點不清楚，也不利于演示。在本節(jié)中，我們將專注于將檢測范圍縮小到溜冰場內(nèi)的物體。此外，我們將把這些邊界框替換為底部的橢圓，以確保更清晰的可見性。
首先，讓我們先從使用方框切換到使用橢圓。為了實現(xiàn)這一點，我們只需在現(xiàn)有代碼中的標簽和邊界框上方添加一個新方法：
#************用橢圓代替邊界框來跟蹤球員的設計**************#
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 frame然后，我們還需要通過調(diào)用橢圓方法替換邊界框和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_frames通過上面這些更改，你的輸出視頻應該看起來更整潔一些，如示例剪輯02所示。
示例剪輯02：用橢圓替換邊界框
現(xiàn)在，為了處理溜冰場的邊界，我們需要對計算機視覺中的分辨率有一些基本的了解。在我們的示例中，我們使用的是720p（1280x720像素）格式；這意味著，我們處理的每一幀或圖像的尺寸為1280像素（寬）乘720像素（高）。
使用720p（1280x720像素）格式意味著什么？這意味著，圖像由水平1280像素和垂直720像素組成。此格式的坐標從圖像左上角的（0,0）開始，x坐標隨著向右移動而增加，y坐標隨著向下移動而增加。這些坐標用于標記圖像中的特定區(qū)域，例如將（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)結果應該是一個綠色矩形，如示例片段03中的（a）所示。但是，為了只跟蹤溜冰場內(nèi)的運動物體，我們需要一個更類似于（b）中的劃界。
圖03：溜冰場的邊界定義
正確地得到（b）中的劃界就像一個反復試驗的過程。在這個過程中，你需要測試不同的坐標，直到找到最適合你的模型的邊界。起初，我的目標是精確地匹配溜冰場的邊界。然而，發(fā)現(xiàn)跟蹤系統(tǒng)在溜冰場的邊緣附近判別比較困難。為了提高準確性，我稍微擴大了一下邊界，以確保捕捉到溜冰場內(nèi)的所有跟蹤對象，同時排除場外的跟蹤對象。最后，如（b）中顯示的結果是我所能得到的最好的結果（你仍然可以在更好的情況下工作），這個邊界由下面這些關鍵坐標定義：
- 左下角：（-450，710）
- 右下角：（2030，710）
- 左上角：（352，61）
- 右上角：（948，61）
最后，我們將定義另外兩個區(qū)域：白隊和黃隊的進攻區(qū)（每支球隊的目標都是得分）。這將使我們能夠收集對手區(qū)域內(nèi)每支球隊的一些基本位置統(tǒng)計數(shù)據(jù)和壓力指標。
圖04：進攻區(qū)
#**************黃隊進攻區(qū)****************
Bottom Left Corner: (-450, 710)
Bottom Right Corner: (2030, 710)
Upper Left Corner: (200, 150)
Upper Right Corner: (1160, 150)
#**************白隊進攻區(qū)****************
Bottom Left Corner: (180, 150)
Bottom Right Corner: (1100, 150)
Upper Left Corner: (352, 61)
Upper Right Corner: (900, 61)我們將暫時擱置這些坐標，并在下一步解釋我們將如何對每個團隊進行分類。然后，我們將把它們整合到我們最初的跟蹤方法中。
使用深度學習進行團隊預測
自Warren McCulloch和Walter Pitts于1943年發(fā)表論文《神經(jīng)活動中內(nèi)在思想的邏輯演算》（https://www.cs.cmu.edu/~./epxing/Class/10715/reading/McCulloch.and.Pitts.pdf）以來，已經(jīng)過去了80多年，該論文為早期神經(jīng)網(wǎng)絡研究奠定了堅實的基礎。后來，在1957年，一個簡化神經(jīng)元的數(shù)學模型（接收輸入，對這些輸入應用權重，求和并輸出二進制結果）啟發(fā)了Frank Rosenblatt構建Mark I（https://news.cornell.edu/stories/2019/09/professors-perceptron-paved-way-ai-60-years-too-soon）。這是第一個旨在演示感知器概念的硬件實現(xiàn)，感知器是一種能夠從數(shù)據(jù)中學習以進行二進制分類的神經(jīng)網(wǎng)絡模型。從那時起，讓計算機像我們一樣思考的探索就沒有放緩。如果這是你第一次深入學習神經(jīng)網(wǎng)絡，或者你想更新和加強你的知識，那么我建議你閱讀Shreya Rao的系列文章（https://medium.com/@shreya.rao/list/deep-learning-illustrated-ae6c27de1640），這可以作為深度學習的一個很好的起點。此外，你可以訪問我鏈接處https://medium.com/@raul.vizcarrach/list/neural-networks-098e9b594f19收集的故事集（來自于不同的投稿者），你可能會發(fā)現(xiàn)這些故事很有用。
那么，為什么選擇卷積神經(jīng)網(wǎng)絡（CNN）呢？老實說，這不是我的第一選擇。最初，我嘗試使用LandingAI構建一個模型，LandingAI是一個用戶友好的云部署平臺，通過API連接Python。然而，這一方案出現(xiàn)了延遲問題（在線處理超過1000幀）。盡管Roboflow中的預訓練模型具有高質(zhì)量的數(shù)據(jù)集和預訓練模型，但它們也出現(xiàn)了類似的延遲問題。意識到需要在本地運行它，我嘗試了一種基于MSE的方法來對球衣顏色進行分類，以便進行球隊和裁判檢測。雖然這聽起來像是最終的解決方案，但它的準確性很低。經(jīng)過幾天的反復試驗，我最終改用CNN。在不同的深度學習方法中，CNN非常適合對象檢測，不像LSTM或RNN更適合語言轉(zhuǎn)錄或翻譯等順序數(shù)據(jù)。
在深入分析代碼之前，讓我們先了解一下CNN架構的一些基本概念：
- 學習樣本數(shù)據(jù)集：數(shù)據(jù)集分為三類：裁判、Team_Away（白色球衣球員）和Team_Home（黃色球衣球員）。每個類別的樣本被分為兩組：訓練數(shù)據(jù)和驗證數(shù)據(jù)。CNN將在每個訓練輪次中使用訓練數(shù)據(jù)來“學習”多層模式。驗證數(shù)據(jù)將在每次迭代結束時用于評估模型的性能，并衡量其對新數(shù)據(jù)的泛化程度。創(chuàng)建樣本數(shù)據(jù)集并不難；我花了大約30到40分鐘從視頻中裁剪每個類別的示例圖像，并將其組織到子目錄中。我設法創(chuàng)建了一個大約90張圖片的示例數(shù)據(jù)集，你可以在項目的GitHub存儲庫中找到。
- 模型是如何學習的？：輸入數(shù)據(jù)在神經(jīng)網(wǎng)絡的每一層中移動，神經(jīng)網(wǎng)絡可以有一層或多層連接在一起進行預測。每一層都使用一個激活函數(shù)來處理數(shù)據(jù)，以進行預測或?qū)?shù)據(jù)進行更改。這些層之間的每個連接都有一個權重，它決定了一個層的輸出對下一個層有多大的影響。目標是找到這些權重的正確組合，以盡量減少預測結果時的錯誤。通過稱為反向傳播和損失函數(shù)的過程，該模型調(diào)整這些權重以減少誤差并提高準確性。這個過程在所謂的訓練輪次（前向傳遞+反向傳播）中重復，隨著模型從錯誤中學習，它在每個周期的預測能力越來越好。
- 激活函數(shù)：如前所述，激活函數(shù)在模型的學習過程中起著重要作用。我選擇了ReLU（校正線性單元）算法，因為它以計算效率高和緩解所謂的消失梯度問題（多層網(wǎng)絡可能會有效地停止學習）而聞名。雖然ReLU工作良好，但其他激活函數(shù)如sigmoid、tanh或swish等也各有其用途，具體取決于網(wǎng)絡的復雜程度。
- 訓練輪次（epoch）：設定正確的訓練輪次需要實驗。你應該考慮數(shù)據(jù)集的復雜性、CNN模型的架構和計算資源等因素。在大多數(shù)情況下，最好在每次迭代中監(jiān)控模型的性能，并在改進變得最小時停止訓練，以防止過擬合?？紤]到我的小訓練數(shù)據(jù)集，我決定以10個訓練輪次作為基線。然而，在其他情況下，可能需要根據(jù)指標性能和驗證結果進行調(diào)整。
- Adam（自適應矩估計）：最終，目標是減少預測輸出和真實輸出之間的誤差。如前所述，反向傳播在這里起著關鍵作用，它通過調(diào)整和更新神經(jīng)網(wǎng)絡權重來隨著時間的推移改進預測。雖然反向傳播基于損失函數(shù)的梯度處理權重更新，但Adam算法通過動態(tài)調(diào)整學習率來逐步最小化誤差或損失函數(shù)，從而增強了這一過程。換句話說，它可以微調(diào)模型的學習速度。
也就是說，為了運行我們的CNN模型，我們需要以下Python包：
pip install torch torchvision
pip install matplotlib
pip install scikit-learn【Tip-02】確保PyTorch安裝正確。我所有的工具都是在Anaconda環(huán)境中設置的，當我安裝PyTorch時，起初它似乎設置得很正確。然而，在運行一些庫時卻出現(xiàn)了一些問題。起初，我以為這是代碼原因，但經(jīng)過幾次修改都沒有成功，最后我不得不重新安裝Anaconda并在干凈的環(huán)境中重新安裝了PyTorch。最終，問題就解決了！
接下來，我們將指定我們的庫和樣本數(shù)據(jù)集的路徑：
# ************卷積神經(jīng)網(wǎng)絡三類檢測**************************
# 裁判
# 白隊(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
#培訓和驗證數(shù)據(jù)集
#從項目的GitHub存儲庫下載teams_sample_dataset文件
data_dir = 'D:/PYTHON/teams_sample_dataset'首先，我們將確保每張圖片大小相等（調(diào)整為150x150像素），然后將其轉(zhuǎn)換為代碼可以理解的格式（在PyTorch中，輸入數(shù)據(jù)通常表示為Tensor對象）。最后，我們將調(diào)整顏色，使模型更容易使用（歸一化），并設置一個加載圖像的過程。這些步驟共同幫助準備圖片并對其進行組織，以便模型能夠有效地開始從中學習，避免數(shù)據(jù)格式引起的偏差。
#******************************數(shù)據(jù)轉(zhuǎn)換***********************************
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])
])
#加載數(shù)據(jù)集
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)接下來，我們將定義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 x我們的CNN模型有三層（conv1、conv2、conv3）。數(shù)據(jù)從卷積層（conv）開始，在那里應用了激活函數(shù)（ReLU）。此函數(shù)使網(wǎng)絡能夠?qū)W習數(shù)據(jù)中的復雜模型和關系。隨后，池化層被激活。什么是最大池化？這是一種在保留重要特征的同時減小圖像大小的技術，有助于高效訓練和優(yōu)化內(nèi)存資源。這個過程在conv1到conv3之間重復。最后，數(shù)據(jù)通過完全連接的層（fc1、fc2）進行最終分類（或決策）。
下一步，我們初始化模型，將分類交叉熵配置為損失函數(shù)（通常用于分類任務），并指定Adam作為我們的優(yōu)化器。如前所述，我們將在10個訓練輪次的完整周期內(nèi)執(zhí)行我們的模型。
#********************************CNN訓練**********************************************
# 模型損失函數(shù)－優(yōu)化器
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())為了跟蹤性能，我們將添加一些代碼來跟蹤訓練進度，打印驗證指標并繪制出來。最后，我們將模型保存為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')此外，在運行完上述所有步驟（你可以在項目的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.9769在完成10個迭代周期后，CNN模型的性能指標有所改善。最初，在第一個訓練輪次中，該模型的訓練損失為1.5346，驗證準確率為47.37%。我們應該如何理解這種初始訓練結果呢？
準確性是評估分類性能的最常見指標之一。在我們的例子中，它代表了正確預測的類別在總數(shù)中所占的比例。然而，僅靠高精度并不能保證整體模型性能；你仍然可能對特定類別的預測得到很差的預測結果（正如我在早期試驗中所經(jīng)歷的那樣）。關于訓練損失函數(shù)，它衡量模型學習將輸入數(shù)據(jù)映射到正確標簽的有效性。由于我們使用的是分類函數(shù)，交叉熵損失量化了預測的類別概率和實際標簽之間的差異。像1.5346這樣的起始值表示預測類別和實際類別之間存在顯著差異；理想情況下，隨著訓練的進行，該值應接近0。隨著時間的推移，我們觀察到訓練損失顯著下降，驗證準確性提高。到最后一個訓練輪次，訓練和驗證損失分別達到0.2509和0.2651的低點。
為了測試我們的CNN模型，我們可以選擇一個球員圖像樣本并評估其預測能力。為了進行測試，你可以運行以下代碼并使用項目GitHub存儲庫中的validation_dataset文件夾。
# *************用樣本數(shù)據(jù)集測試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
# 用于驗證的樣本數(shù)據(jù)集
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):
# 加載數(shù)據(jù)集
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}')輸出應該看起來像下面這樣：
# *************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_referee正如你所看到的，該模型在識別球隊和排除裁判作為團隊球員方面表現(xiàn)出了很好的能力。
【提示03】我在CNN設計過程中學到的一點是，增加復雜性并不總是能提高性能。最初，我嘗試了更深層次的模型（更多的卷積層）和基于顏色的增強，以提高球員的球衣識別率。然而，在我的小數(shù)據(jù)集中，我遇到了過擬合，而不是學習可泛化的特征（所有圖像都被預測為白人球隊球員或裁判）。像dropout和批歸一化這樣的正則化技術也很重要；它們有助于在訓練過程中施加約束，確保模型能夠很好地泛化到新數(shù)據(jù)。就結果而言，更少有時意味著更多。
組合應用
將所有上面這些技術放在一起應用時需要對前面描述的跟蹤機制進行一些調(diào)整。下面描述更新代碼的逐步分解過程。
首先，我們將設置所需的庫和路徑。請注意，現(xiàn)在指定了pickle文件和CNN模型的路徑。這一次，如果在路徑中找不到pickle文件，代碼將拋出錯誤。如果需要，使用前面的代碼生成pickle文件，并使用此更新版本執(zhí)行視頻分析：
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'接下來，我們將加載模型，指定溜冰場坐標，并像以前一樣，以20個為一批啟動檢測每幀中對象的過程。請注意，目前，我們將只使用溜冰場邊界來將分析重點放在溜冰場上。在本文的最后步驟中，當我們包含性能統(tǒng)計數(shù)據(jù)時，我們將使用進攻區(qū)域坐標。
#***************************加載模型和溜冰場坐標********************#
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 detections接下來，我們將添加預測每個球員球隊的過程：
#*********************** 加載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 team下一步，我們將添加前面描述的方法，從邊界框切換到橢圓：
#************ 使用橢圓來跟蹤球員，而不是邊界框*******************#
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 frame現(xiàn)在，是時候添加分析器了，包括讀取pickle文件，將分析范圍縮小到我們之前定義的溜冰場邊界內(nèi)，并調(diào)用CNN模型來識別每個球員的球隊并添加標簽。請注意，我們提供了一個函數(shù)，可以用不同的顏色標記裁判，并更改其橢圓的顏色。在代碼的最后部分，將處理后的幀寫入輸出視頻。
#*******************加載跟蹤的數(shù)據(jù) (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
#***********檢查球員是否落入溜冰場區(qū)域**********************************#
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"]
# 檢查球員是否落入溜冰場區(qū)域
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()最后，我們實現(xiàn)CNN的架構（在CNN設計過程中定義）并執(zhí)行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
#*********執(zhí)行HockeyAnalyzer/分類器并保存輸出***********#
analyzer = HockeyAnalyzer(model_path, classifier_path)
analyzer.analyze_video(video_path, output_path, tracks_path)運行所有步驟后，你的視頻輸出應該如下所示：
示例剪輯06：跟蹤球員和團隊
請注意，在最后一次更新中，物體檢測僅在溜冰場內(nèi)進行，球隊和裁判都是不同的。雖然CNN模型仍需要微調(diào)，偶爾會對一些球員失去穩(wěn)定性，但它在整個視頻中仍然基本可靠和準確。
速度、距離和進攻壓力
實驗結果證明，本示例開發(fā)的模型方案在跟蹤球隊和球員信息方面展現(xiàn)出強大能力，這為衡量各種指標的應用開辟了令人興奮的可能性，例如生成熱圖、分析速度和覆蓋距離、跟蹤區(qū)域入口或出口等動作，甚至還可以更深入地了解詳細的球員指標。不妨讓我們作一個測試，我們添加三個性能指標：每位球員的平均速度、每支球隊的滑冰距離和進攻壓力（以每支球隊在對手區(qū)域內(nèi)投入的距離值百分比來衡量）。我會把更詳細的統(tǒng)計數(shù)據(jù)留給你！
最開始，我們將溜冰場的坐標從基于像素的測量值調(diào)整為基于米單位。這種調(diào)整使我們能夠以米而不是像素讀取數(shù)據(jù)。視頻中看到的溜冰場的真實尺寸約為15mx30m（寬15米，高30米）。為了便于這種轉(zhuǎn)換，我們引入了一種將像素坐標轉(zhuǎn)換為米的方法。通過定義溜冰場的實際尺寸并使用其角的像素坐標（從左到右和從上到下），我們獲得了轉(zhuǎn)換因子。這些因素將支持我們估算以米為單位的距離和以米每秒為單位的速度的過程。（你可以探索和應用的另一種有趣的技術是透視變換）
#*********************加載模型和溜冰場坐標*****************#
class_names = ['Referee', 'Tm_white', 'Tm_yellow']
class HockeyAnalyzer:
def __init__(self, model_path, classifier_path):
*
*
*
*
*
*
self.pixel_to_meter_conversion() #<------ 添加這個工具方法
#***********將溜冰場的坐標從基于像素的測量值調(diào)整為基于米單位***************************#
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]
#轉(zhuǎn)換因子
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):
#把像素轉(zhuǎn)換成米
return distance_pixels / self.pixels_per_meter_x, distance_pixels / self我們現(xiàn)在準備為每個球員增加速度，以米每秒為單位。為此，我們需要進行三處修改。首先，在HockeyAnalyzer類中啟動一個名為previous_positions的空字典，以幫助我們比較球員的當前和先前位置。同樣，我們將創(chuàng)建一個team_stats結構來存儲每個團隊的統(tǒng)計數(shù)據(jù)，以便進一步可視化。
接下來，我們將添加一種速度方法來估計球員的速度（以每秒像素為單位），然后使用轉(zhuǎn)換因子（前面解釋過）將其轉(zhuǎn)換為每秒米數(shù)。最后，根據(jù)analyze_video方法，我們將調(diào)用新的速度方法，并將速度添加到每個被跟蹤的對象（球員和裁判）中。下面給出的是修改后的代碼：
#*********************加載模型和溜冰場坐標*****************#
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
#******************* 加載跟蹤數(shù)據(jù) (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()如果你在添加上面這幾行新代碼時遇到問題，你可以隨時訪問該項目的GitHub存儲庫（https://github.com/rvizcarra15/IceHockey_ComputerVision_PyTorch），在那里你可以找到完整的參考代碼。此時，你的視頻輸出結果應該像下圖的樣子（請注意，速度已添加到每個播放器的標簽中）：
示例剪輯07：跟蹤球員和速度
最后，讓我們添加一個統(tǒng)計板，借助這個統(tǒng)計板我們可以跟蹤每支球隊每位球員的平均速度，以及其他指標，如覆蓋距離和對手區(qū)域的進攻壓力。
我們已經(jīng)定義了進攻區(qū)，并將其整合到我們的代碼中?，F(xiàn)在，我們需要跟蹤每個球員進入對手區(qū)域的頻率。為了實現(xiàn)這一點，我們將使用光線投射算法實現(xiàn)一種方法。該算法檢查球員的位置是否在白隊或黃隊的進攻區(qū)域內(nèi)。它的工作原理是從球員到目標區(qū)域畫一條假想線。如果線穿過一個邊界，則球員在里面；如果它穿過更多邊界（在我們的例子中，四個邊界中的兩個），則球員就在外面。然后，該代碼掃描整個視頻，以確定每個被跟蹤對象的區(qū)域狀態(tài)。
#************ 在目標區(qū)域中定位球員的位置***********************#
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 inside現(xiàn)在，我們將通過添加一種方法來處理表現(xiàn)指標，該方法以表格形式顯示每支球隊的平均球員速度、總覆蓋距離和進攻壓力（在對手區(qū)域花費的時間百分比）。使用OpenCV，我們將把這些指標格式化為覆蓋在視頻上的表格，并引入動態(tài)更新機制，以便在游戲過程中保持實時統(tǒng)計數(shù)據(jù)。
#*******************************性能度量*********************************************#
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) # 在此處調(diào)整行距
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,
)
#****************** 跟蹤和更新游戲統(tǒng)計數(shù)據(jù)****************************************#
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
為了在視頻中顯示統(tǒng)計數(shù)據(jù)，我們必須調(diào)用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 += 1每個球員覆蓋的距離（以米為單位）是通過將他們的速度（以米每秒為單位）除以幀率（每秒幀數(shù)）來計算的。這種計算使我們能夠估計每個球員在視頻中的每一幀變化之間移動了多遠。如果一切正常，你的最終視頻輸出應該如下：
示例剪輯08：示例程序最終輸出結果
待考慮因素和未來的工作
本文示例中實現(xiàn)的模型是使用計算機視覺跟蹤冰球比賽（或任何團隊運動）中有關球員的基本設置。但是，可以進行很多微調(diào)來改進它并添加新功能。以下是我為下一個2.0版本所做的一些想法，你也可以考慮：
- 跟蹤冰球的挑戰(zhàn)：考慮到冰球與足球或籃球相比的大小，跟蹤冰球具有挑戰(zhàn)性，這取決于你的相機朝向哪個方向和分辨率。但如果你實現(xiàn)了這一點，那么跟蹤將表現(xiàn)出非常有趣的可能性，比如控球時間指標、進球機會或投籃數(shù)據(jù)。這也適用于個人表演；在冰球中，球員的變化比其他團隊運動要頻繁得多，因此跟蹤每個球員在一個時期的表現(xiàn)是一個挑戰(zhàn)。
- 計算資源，哦，為什么要計算呢？我在CPU上運行了所有代碼，但遇到了問題（有時會導致藍屏），由于在設計過程中內(nèi)存不足（考慮使用CUDA設置）。我們的示例視頻長約40秒，最初大小為5MB，但在運行模型后，輸出增加到34MB。想象一下一個20分
- 不要低估MLOps：為了快速部署和擴展，我們需要高效、支持頻繁執(zhí)行且可靠的機器學習管道。這涉及到考慮采用持續(xù)集成部署訓練方法。我們的用例是為特定場景構建的，但如果條件發(fā)生變化，比如相機方向或球衣顏色，該怎么辦？為了擴大規(guī)模，我們必須采取CI/CD/CT的思維方式。
最后，我希望你覺得本文提供的這個計算機視覺演示項目非常有趣，你可以在GitHub存儲庫（https://github.com/rvizcarra15/IceHockey_ComputerVision_PyTorch）中訪問完整的代碼。如果你想支持該地區(qū)直排冰球和冰球的發(fā)展，請關注APHL（我們總是需要你想為年輕球員捐贈的二手設備，并致力于建設我們的第一個官方冰球場），在全球范圍內(nèi)，請關注并支持友誼聯(lián)盟（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。
