Spaces:
Sleeping
Sleeping
| import cv2 | |
| import gradio as gr | |
| import tempfile | |
| from ultralytics import YOLO | |
| from deep_sort_realtime.deepsort_tracker import DeepSort | |
| import os | |
| # Load model and tracker | |
| model = YOLO("yolov8n.pt") | |
| tracker = DeepSort(max_age=30) | |
| def track_players(video): | |
| # Create temp output path | |
| temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) | |
| output_path = temp_out.name | |
| cap = cv2.VideoCapture(video) | |
| if not cap.isOpened(): | |
| raise IOError("Video file could not be opened.") | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 | |
| out = cv2.VideoWriter( | |
| output_path, | |
| cv2.VideoWriter_fourcc(*"avc1"), # 'avc1' for Linux compatibility | |
| fps, | |
| (width, height) | |
| ) | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| results = model(frame)[0] | |
| detections = [] | |
| for result in results.boxes.data.tolist(): | |
| x1, y1, x2, y2, score, class_id = result | |
| if int(class_id) == 0 and score > 0.4: | |
| detections.append(([x1, y1, x2 - x1, y2 - y1], score, 'player')) | |
| tracks = tracker.update_tracks(detections, frame=frame) | |
| for track in tracks: | |
| if not track.is_confirmed(): | |
| continue | |
| track_id = track.track_id | |
| l, t, r, b = track.to_ltrb() | |
| cv2.rectangle(frame, (int(l), int(t)), (int(r), int(b)), (0, 255, 0), 2) | |
| cv2.putText(frame, f'Player {track_id}', (int(l), int(t - 10)), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) | |
| out.write(frame) | |
| cap.release() | |
| out.release() | |
| return output_path | |
| demo = gr.Interface( | |
| fn=track_players, | |
| inputs=gr.Video(label="Upload Match Video"), | |
| outputs=gr.Video(label="Tracked Output"), | |
| title="🏟️ Football Player Tracker", | |
| description="Upload a match video and get players tracked with consistent IDs." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |