Spaces:
Runtime error
Runtime error
File size: 3,661 Bytes
e19f940 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
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()
|