Finetuned YOLO and Supervision Basketball Video Tracker

Automatically track players and the ball in NBA broadcast videos using computer vision.

video_1.mp4

What It Does

• Detects and tracks all players on the court
• Tracks the basketball throughout the game
• Assigns players to their teams based on uniform colors
• Detects who has possession of the ball
• Outputs an annotated video with visual overlays

How It Works

1. Detection: Uses YOLOv5 models to find players and the ball in each frame
2. Tracking: Follows detected objects across frames with ByteTrack (via Supervision)
3. Team Assignment: Groups players by uniform color
4. Ball Possession: Determines which player controls the ball
5. Visualization: Draws bounding boxes and labels on the video

Configuration

Edit main.py to change:

• Input video: video_path="input_videos/your_video.mp4"
• Output location: output_video_path="output_videos/output.avi"
• Caching: Set read_from_stub=False to reprocess from scratch

Models & Tools

Detection Models - Fine-tuned YOLOv5:

• player_detection.pt - Detects players
• ball_detection.pt - Detects the basketball

Tracking - Supervision library with ByteTrack algorithm for multi-object tracking

Training Data - Roboflow Universe