大数据文摘授权转载自法纳斯特
作者：小F
大家好～
今天给大家介绍一个计算机视觉实战的项目。
该项目使用YOLO算法检测球员和网球，并利用cnn提取球场关键点。
进而分析视频中的网球运动员，测量他们的速度、击球速度和击球次数。
使用win10电脑，Python 3.9.7，相关依赖版本如下。
numpy==1.22.4opencv_python==4.8.0.74pandas==2.2.2torch==2.0.1torchvision==0.15.2ultralytics==8.0.178
可以使用conda创建Python环境，然后执行主程序。
电脑无需GPU，普通CPU电脑即可~
# 创建Python环境conda create --name tennis_analysis python=3.9.7# 激活环境conda activate tennis_analysis# 安装依赖pip install -r requirements.txt -i 网页链接 执行程序python main.py
主程序代码如下。
from utils import (read_video, save_video, measure_distance, draw_player_stats, convert_pixel_distance_to_meters )import constantsfrom trackers import PlayerTracker, BallTrackerfrom court_line_detector import CourtLineDetectorfrom mini_court import MiniCourtimport cv2import pandas as pdfrom copy import deepcopy
def main(): # Read Video input_video_path = "input_videos/input_video.mp4" video_frames = read_video(input_video_path)
# Detect Players and Ball player_tracker = PlayerTracker(model_path='yolov8x') ball_tracker = BallTracker(model_path='models/yolo5_last.pt')
player_detections = player_tracker.detect_frames(video_frames, read_from_stub=True, stub_path="tracker_stubs/player_detections.pkl" ) ball_detections = ball_tracker.detect_frames(video_frames, read_from_stub=True, stub_path="tracker_stubs/ball_detections.pkl" ) ball_detections = ball_tracker.interpolate_ball_positions(ball_detections)
# Court Line Detector model court_model_path = "models/keypoints_model.pth" court_line_detector = CourtLineDetector(court_model_path) court_keypoints = court_line_detector.predict(video_frames[0])
# choose players player_detections = player_tracker.choose_and_filter_players(court_keypoints, player_detections)
# MiniCourt mini_court = MiniCourt(video_frames[0])
# Detect ball shots ball_shot_frames = ball_tracker.get_ball_shot_frames(ball_detections)
# Convert positions to mini court positions player_mini_court_detections, ball_mini_court_detections = mini_court.convert_bounding_boxes_to_mini_court_coordinates( player_detections, ball_detections, court_keypoints)
player_stats_data = [{ 'frame_num': 0, 'player_1_number_of_shots': 0, 'player_1_total_shot_speed': 0, 'player_1_last_shot_speed': 0, 'player_1_total_player_speed': 0, 'player_1_last_player_speed': 0,
'player_2_number_of_shots': 0, 'player_2_total_shot_speed': 0, 'player_2_last_shot_speed': 0, 'player_2_total_player_speed': 0, 'player_2_last_player_speed': 0, }]
for ball_shot_ind in range(len(ball_shot_frames) - 1): start_frame = ball_shot_frames[ball_shot_ind] end_frame = ball_shot_frames[ball_shot_ind + 1] ball_shot_time_in_seconds = (end_frame - start_frame) / 24 # 24fps
# Get distance covered by the ball distance_covered_by_ball_pixels = measure_distance(ball_mini_court_detections[start_frame][1], ball_mini_court_detections[end_frame][1]) distance_covered_by_ball_meters = convert_pixel_distance_to_meters(distance_covered_by_ball_pixels, constants.DOUBLE_LINE_WIDTH, mini_court.get_width_of_mini_court() )
# Speed of the ball shot in km/h speed_of_ball_shot = distance_covered_by_ball_meters / ball_shot_time_in_seconds * 3.6
# player who the ball player_positions = player_mini_court_detections[start_frame] player_shot_ball = min(player_positions.keys(), key=lambda player_id: measure_distance(player_positions[player_id], ball_mini_court_detections[start_frame][1]))
# opponent player speed opponent_player_id = 1 if player_shot_ball == 2 else 2 distance_covered_by_opponent_pixels = measure_distance( player_mini_court_detections[start_frame][opponent_player_id], player_mini_court_detections[end_frame][opponent_player_id]) distance_covered_by_opponent_meters = convert_pixel_distance_to_meters(distance_covered_by_opponent_pixels, constants.DOUBLE_LINE_WIDTH, mini_court.get_width_of_mini_court() )
speed_of_opponent = distance_covered_by_opponent_meters / ball_shot_time_in_seconds * 3.6
current_player_stats = deepcopy(player_stats_data[-1]) current_player_stats['frame_num'] = start_frame current_player_stats[f'player_{player_shot_ball}_number_of_shots'] += 1 current_player_stats[f'player_{player_shot_ball}_total_shot_speed'] += speed_of_ball_shot current_player_stats[f'player_{player_shot_ball}_last_shot_speed'] = speed_of_ball_shot
current_player_stats[f'player_{opponent_player_id}_total_player_speed'] += speed_of_opponent current_player_stats[f'player_{opponent_player_id}_last_player_speed'] = speed_of_opponent
player_stats_data.append(current_player_stats)
player_stats_data_df = pd.DataFrame(player_stats_data) frames_df = pd.DataFrame({'frame_num': list(range(len(video_frames)))}) player_stats_data_df = pd.merge(frames_df, player_stats_data_df, on='frame_num', how='left') player_stats_data_df = player_stats_data_df.ffill()
player_stats_data_df['player_1_average_shot_speed'] = player_stats_data_df['player_1_total_shot_speed'] / \ player_stats_data_df['player_1_number_of_shots'] player_stats_data_df['player_2_average_shot_speed'] = player_stats_data_df['player_2_total_shot_speed'] / \ player_stats_data_df['player_2_number_of_shots'] player_stats_data_df['player_1_average_player_speed'] = player_stats_data_df['player_1_total_player_speed'] / \ player_stats_data_df['player_2_number_of_shots'] player_stats_data_df['player_2_average_player_speed'] = player_stats_data_df['player_2_total_player_speed'] / \ player_stats_data_df['player_1_number_of_shots']
# Draw output ## Draw Player Bounding Boxes output_video_frames = player_tracker.draw_bboxes(video_frames, player_detections) output_video_frames = ball_tracker.draw_bboxes(output_video_frames, ball_detections)
## Draw court Keypoints output_video_frames = court_line_detector.draw_keypoints_on_video(output_video_frames, court_keypoints)
# Draw Mini Court output_video_frames = mini_court.draw_mini_court(output_video_frames) output_video_frames = mini_court.draw_points_on_mini_court(output_video_frames, player_mini_court_detections) output_video_frames = mini_court.draw_points_on_mini_court(output_video_frames, ball_mini_court_detections, color=(0, 255, 255))
# Draw Player Stats output_video_frames = draw_player_stats(output_video_frames, player_stats_data_df)
## Draw frame number on top left corner for i, frame in enumerate(output_video_frames): cv2.putText(frame, f"Frame: {i}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
save_video(output_video_frames, "output_videos/output_video.avi")
if __name__ == "__main__": main()
最终可以在output_videos文件夹下看到结果。
结果如下所示。
这里还提供了网球模型训练的代码，大家可以使用Colab或Kaggle的免费GPU进行训练。
感兴趣的小伙伴，可以自行去学习下~
租售GPU算力
租：4090/A800/H800/H100
售：现货H100/H800
特别适合企业级应用
扫码了解详情☝
点「在看」的人都变好看了哦！