Spaces:
Runtime error
Runtime error
| from click import launch | |
| from player_ball_assigner import PlayerBallAssigner | |
| from utils import read_video,save_video,check_video_resolution | |
| from trackers import Tracker | |
| from team_assigner import TeamAssigner | |
| import numpy as np | |
| from camera_movement_estimator import CameraMovementEstimator | |
| from view_transformer import ViewTransformer | |
| from speed_and_distance_estimator import SpeedAndDistance_Estimator | |
| import gradio as gr | |
| def annotate(input): | |
| # Read Video | |
| video_frames = read_video(input) | |
| #intialize tracker | |
| tracker = Tracker('models/best.pt') | |
| tracks = tracker.get_object_tracks(video_frames,read_from_stub=False,stub_path=None) | |
| # Get Object Postions | |
| tracker.add_postion_to_tracks(tracks) | |
| # camera movement estimator | |
| camera_movement_estimator = CameraMovementEstimator(video_frames[0]) | |
| camera_movement_per_frame = camera_movement_estimator.get_camera_movement(video_frames, | |
| read_from_stub=False, | |
| stub_path=None) | |
| camera_movement_estimator.add_adjust_positions_to_tracks(tracks,camera_movement_per_frame) | |
| # View Trasnformer | |
| view_transformer = ViewTransformer() | |
| view_transformer.add_transformed_position_to_tracks(tracks) | |
| # interpolate ball positions | |
| tracks['ball'] = tracker.interpolate_ball_positions(tracks['ball']) | |
| # Speed and distance estimator | |
| speed_and_distance_estimator = SpeedAndDistance_Estimator() | |
| speed_and_distance_estimator.add_speed_and_distance_to_tracks(tracks) | |
| # Assign player teams | |
| team_assigner = TeamAssigner() | |
| team_assigner.assign_team_color(video_frames[0],tracks['players'][0]) | |
| for frame_num, player_track in enumerate(tracks['players']): | |
| for player_id,track in player_track.items(): | |
| team = team_assigner.get_player_team(video_frames[frame_num], | |
| track['bbox'], | |
| player_id) | |
| tracks['players'][frame_num][player_id]['team'] = team | |
| tracks['players'][frame_num][player_id]['team_color'] = team_assigner.team_colors[team] | |
| # Assign Ball Acquisition | |
| player_assigner = PlayerBallAssigner() | |
| team_ball_control = [] | |
| for frame_num,player_track in enumerate(tracks['players']): | |
| ball_bbox = tracks['ball'][frame_num][1]['bbox'] | |
| assigned_player = player_assigner.assign_ball_to_player(player_track,ball_bbox) | |
| if assigned_player != -1: | |
| tracks['players'][frame_num][assigned_player]['has_ball'] = True | |
| team_ball_control.append(tracks['players'][frame_num][assigned_player]['team']) | |
| else: | |
| team_ball_control.append(team_ball_control[-1]) | |
| team_ball_control = np.array(team_ball_control) | |
| # Draw Output | |
| ## Draw object tracks | |
| output_video_frames = tracker.draw_annotations(video_frames,tracks,team_ball_control) | |
| ## Draw Camera movement | |
| output_video_frames = camera_movement_estimator.draw_camera_movement(output_video_frames,camera_movement_per_frame) | |
| ## Draw Speed and Distance | |
| speed_and_distance_estimator.draw_speed_and_distance(output_video_frames,tracks) | |
| output_path= 'output_videos/out.mp4' | |
| # Save Video | |
| save_video(output_video_frames,output_path) | |
| return output_path | |
| iface = gr.Interface( | |
| fn=annotate, | |
| inputs=['video'], | |
| outputs=['video'], | |
| title="Football Analysis", | |
| examples=['input_videos/08fd33_4.mp4'], | |
| ).queue(default_concurrency_limit=2).launch() | |