基于YOLO算法实现网球运动实时分析(附源码)

大家好，我是小F～

今天给大家介绍一个计算机视觉实战的项目。

该项目使用YOLO算法检测球员和网球，并利用cnn提取球场关键点。

进而分析视频中的网球运动员，测量他们的速度、击球速度和击球次数。

使用win10电脑，Python 3.9.7，相关依赖版本如下。

numpy==1.22.4
opencv_python==4.8.0.74
pandas==2.2.2
torch==2.0.1
torchvision==0.15.2
ultralytics==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 https://mirror.baidu.com/pypi/simple# 执行程序
python main.py

主程序代码如下。

from utils import (read_video,save_video,measure_distance,draw_player_stats,convert_pixel_distance_to_meters)
import constants
from trackers import PlayerTracker, BallTracker
from court_line_detector import CourtLineDetector
from mini_court import MiniCourt
import cv2
import pandas as pd
from copy import deepcopydef main():# Read Videoinput_video_path = "input_videos/input_video.mp4"video_frames = read_video(input_video_path)# Detect Players and Ballplayer_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 modelcourt_model_path = "models/keypoints_model.pth"court_line_detector = CourtLineDetector(court_model_path)court_keypoints = court_line_detector.predict(video_frames[0])# choose playersplayer_detections = player_tracker.choose_and_filter_players(court_keypoints, player_detections)# MiniCourtmini_court = MiniCourt(video_frames[0])# Detect ball shotsball_shot_frames = ball_tracker.get_ball_shot_frames(ball_detections)# Convert positions to mini court positionsplayer_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 balldistance_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/hspeed_of_ball_shot = distance_covered_by_ball_meters / ball_shot_time_in_seconds * 3.6# player who the ballplayer_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 speedopponent_player_id = 1 if player_shot_ball == 2 else 2distance_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.6current_player_stats = deepcopy(player_stats_data[-1])current_player_stats['frame_num'] = start_framecurrent_player_stats[f'player_{player_shot_ball}_number_of_shots'] += 1current_player_stats[f'player_{player_shot_ball}_total_shot_speed'] += speed_of_ball_shotcurrent_player_stats[f'player_{player_shot_ball}_last_shot_speed'] = speed_of_ball_shotcurrent_player_stats[f'player_{opponent_player_id}_total_player_speed'] += speed_of_opponentcurrent_player_stats[f'player_{opponent_player_id}_last_player_speed'] = speed_of_opponentplayer_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 Boxesoutput_video_frames = player_tracker.draw_bboxes(video_frames, player_detections)output_video_frames = ball_tracker.draw_bboxes(output_video_frames, ball_detections)## Draw court Keypointsoutput_video_frames = court_line_detector.draw_keypoints_on_video(output_video_frames, court_keypoints)# Draw Mini Courtoutput_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 Statsoutput_video_frames = draw_player_stats(output_video_frames, player_stats_data_df)## Draw frame number on top left cornerfor 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文件夹下看到结果。