大家好，我是小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文件夾下看到結果。

結果如下所示。

這里還提供了網球模型訓練的代碼，大家可以使用Colab或Kaggle的免費GPU進行訓練。

感興趣的小伙伴，可以自行去學習下~

項目源碼，公眾號后臺回復：「網球分析」，即可獲得。

萬水千山總是情，點個?👍?行不行。

推薦閱讀

···? END? ···