Spaces:
Running
Running
| # app.py | |
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| import supervision as sv | |
| from scipy.spatial import distance as dist | |
| import tempfile | |
| # Initialize components | |
| model = YOLO('yolov8s.pt') | |
| tracker = sv.ByteTrack() | |
| # Streamlit UI | |
| st.title("⚽️ Player Tracking System") | |
| uploaded_video = st.file_uploader("Upload match video", type=["mp4", "mov"]) | |
| calibration_dist = st.number_input("Field width in meters (for speed calibration):", value=68.0) | |
| # Initialize session state | |
| if 'player_data' not in st.session_state: | |
| st.session_state.player_data = {} | |
| if uploaded_video: | |
| tfile = tempfile.NamedTemporaryFile(delete=False) | |
| tfile.write(uploaded_video.read()) | |
| cap = cv2.VideoCapture(tfile.name) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| pixels_per_meter = frame_width / calibration_dist if calibration_dist > 0 else 1.0 | |
| st_frame = st.empty() | |
| frame_count = 0 | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Detection + tracking | |
| results = model(frame)[0] | |
| detections = sv.Detections.from_ultralytics(results) | |
| detections = tracker.update_with_detections(detections) | |
| for detection in detections: | |
| track_id = int(detection[4]) | |
| bbox = detection[0] | |
| centroid = (int((bbox[0]+bbox[2])/2), int((bbox[1]+bbox[3])/2)) | |
| speed = 0.0 # Initialize speed for all cases | |
| # Initialize new player | |
| if track_id not in st.session_state.player_data: | |
| st.session_state.player_data[track_id] = { | |
| 'positions': [centroid], | |
| 'timestamps': [frame_count/fps], | |
| 'distance': 0.0, | |
| 'speeds': [] | |
| } | |
| else: | |
| # Calculate movement metrics | |
| prev_pos = st.session_state.player_data[track_id]['positions'][-1] | |
| time_diff = (frame_count/fps) - st.session_state.player_data[track_id]['timestamps'][-1] | |
| pixel_dist = dist.euclidean(prev_pos, centroid) | |
| speed_px = pixel_dist / time_diff if time_diff > 0 else 0.0 | |
| speed = speed_px * pixels_per_meter # Convert to m/s | |
| # Update player record | |
| st.session_state.player_data[track_id]['positions'].append(centroid) | |
| st.session_state.player_data[track_id]['timestamps'].append(frame_count/fps) | |
| st.session_state.player_data[track_id]['distance'] += pixel_dist | |
| st.session_state.player_data[track_id]['speeds'].append(speed) | |
| # Annotation with failsafe | |
| label = f"ID:{track_id} Speed:{speed:.1f}m/s" | |
| cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), | |
| (int(bbox[2]), int(bbox[3])), (0,255,0), 2) | |
| cv2.putText(frame, label, (int(bbox[0]), int(bbox[1]-10)), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1) | |
| st_frame.image(frame, channels="BGR") | |
| frame_count += 1 | |
| cap.release() | |
| # Display final analytics | |
| st.subheader("Player Statistics") | |
| for player_id, data in st.session_state.player_data.items(): | |
| avg_speed = np.mean(data['speeds']) if data['speeds'] else 0 | |
| st.write(f""" | |
| **Player {player_id}** | |
| - Total distance: {data['distance'] * pixels_per_meter:.2f}m | |
| - Avg speed: {avg_speed:.1f}m/s | |
| - Max speed: {np.max(data['speeds']):.1f}m/s | |
| - Active duration: {data['timestamps'][-1]-data['timestamps'][0]:.1f}s | |
| """) | |