Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import tempfile | |
| import cv2 | |
| from utils import read_video, save_video | |
| from trackers import Tracker | |
| from team_assigner import TeamAssigner | |
| from player_ball_assigner import PlayerBallAssigner | |
| from camera_movment_estimator import CameraMovmentEstimator | |
| from view_transformer import ViewTransformer | |
| from speed_and_distance_estimator import SpeedAndDistance_Estimator | |
| import numpy as np | |
| import os | |
| def main(): | |
| st.title("Football Match Analysis App") | |
| uploaded_video = st.file_uploader("Upload a football clip", type=["mp4", "avi", "mov"]) | |
| if uploaded_video is not None: | |
| # Save the uploaded file temporarily | |
| temp_file = tempfile.NamedTemporaryFile(delete=False) | |
| temp_file.write(uploaded_video.read()) | |
| temp_file_path = temp_file.name | |
| # Read video frames using the uploaded file path | |
| video_frames = read_video(temp_file_path) | |
| tracker = Tracker("./best.pt") | |
| tracks = tracker.get_object_tracks(video_frames, read_from_stub=True, stub_path="./stubs/track_stub.pkl") | |
| camera_movment_estimator = CameraMovmentEstimator(video_frames[0]) | |
| camera_movment_per_frame = camera_movment_estimator.get_camera_movment( | |
| video_frames, read_from_stub=True, stub_path="./stubs/camera_movment_stub.pkl" | |
| ) | |
| camera_movment_estimator.add_position_to_tracks(tracks) | |
| camera_movment_estimator.add_adjust_positions_to_tracks(tracks, camera_movment_per_frame) | |
| view_transformer = ViewTransformer() | |
| view_transformer.add_transformed_position_to_tracks(tracks) | |
| tracks["ball"] = tracker.interppolate_ball_positions(tracks["ball"]) | |
| speed_and_distance_estimator = SpeedAndDistance_Estimator() | |
| speed_and_distance_estimator.add_speed_and_distance_to_tracks(tracks) | |
| 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] | |
| player_ball_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_ball_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) | |
| output_video_frames = tracker.draw_annotations(video_frames, tracks, team_ball_control) | |
| output_video_frames = camera_movment_estimator.draw_camera_movment(output_video_frames, camera_movment_per_frame) | |
| speed_and_distance_estimator.draw_speed_and_distance(output_video_frames, tracks) | |
| # Save the video temporarily in memory for download | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_output: | |
| save_video(output_video_frames, temp_output.name) | |
| temp_output_path = temp_output.name | |
| # Allow user to download the video | |
| with open(temp_output_path, 'rb') as video_file: | |
| video_bytes = video_file.read() | |
| st.download_button( | |
| label="Download Output Video", | |
| data=video_bytes, | |
| file_name="output_video.mp4", | |
| mime="video/mp4" | |
| ) | |
| # Clean up the temporary file after displaying | |
| os.remove(temp_output_path) | |
| if __name__ == "__main__": | |
| main() |