-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
113 lines (89 loc) · 3.7 KB
/
main.py
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
104
105
106
107
108
109
110
111
112
113
from utils import read_video, save_video
import cv2
import numpy as np
from trackers import Tracker
from team_assigner import TeamAssigner
from player_ball_assigner import PlayerBallAssigner
from camera_movement_estimator import CameraMovementEstimator
from view_transformer import ViewTransformer
from speed_and_distance_estimator import SpeedAndDistanceEstimator
def main():
# Read video
video_frames = read_video("input_videos/08fd33_4.mp4")
# Initialize tracker
tracker = Tracker("models/best.pt")
tracks = tracker.get_object_tracks(
video_frames, read_from_stub=True, stub_path="stubs/track_stubs.pkl"
)
# Get object positions
tracker.add_position_to_track(tracks)
"""
# Save cropped image of a player
for track_id, player in tracks["players"][0].items():
bbox = player["bbox"]
frame = video_frames[0]
# Crop bbox from frame
cropped_image = frame[int(bbox[1]) : int(bbox[3]), int(bbox[0]) : int(bbox[2])]
# Save the cropped image
cv2.imwrite(f"output_videos/cropped_image.jpg", cropped_image)
break
"""
# Estimate camera movement
camera_movement_estimator = CameraMovementEstimator(video_frames[0])
camera_movement_per_frame = camera_movement_estimator.get_camera_movement(
video_frames, read_from_stub=True, stub_path="stubs/camera_movement_stubs.pkl"
)
camera_movement_estimator.add_adjust_positions_to_tracks(
tracks, camera_movement_per_frame
)
# View Transformer
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 = SpeedAndDistanceEstimator()
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_tracks in enumerate(tracks["players"]):
for player_id, player in player_tracks.items():
player_team = team_assigner.get_player_team(
video_frames[frame_num], player["bbox"], player_id
)
tracks["players"][frame_num][player_id]["team"] = player_team
tracks["players"][frame_num][player_id]["team_color"] = (
team_assigner.team_colors[player_team]
)
# Assign ball to player
player_assigner = PlayerBallAssigner()
team_ball_control = []
for frame_num, player_tracks in enumerate(tracks["players"]):
ball_bbox = tracks["ball"][frame_num][1]["bbox"]
assigned_player = player_assigner.assign_ball_to_player(
player_tracks, 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)
# Save video
save_video(output_video_frames, "output_videos/output_video.avi")
if __name__ == "__main__":
main()