This project uses machine learning, computer vision, and deep learning techniques to create a football analysis system. The main goals are:
- Detect and track players, referees, and the ball across video frames using object detection and tracking algorithms.
- Train a custom object detector to enhance the output of state-of-the-art models.
- Assign players to teams based on the color of their jerseys using K-means clustering for pixel segmentation.
- Analyze ball possession and control by assigning the ball to the closest player.
- Annotate output video with player IDs, teams, ball possession, and other relevant information.
- Clone the repository
- Install required dependencies:
pip install -r requirements.txt
- Train your model using the training notebook provided
- Download the model pt file and store in models folder
- Store the video to be analysed in the input folder
To run the analysis pipeline:
python main.py
This will process the specified input video, perform analysis, and generate an annotated output video in the output folder. If the video has been processed before, the code will use the previously generated pickle file to save time, instead of reprocessing the video from scratch.
main.py
: The main entry point that orchestrates the analysis pipeline.utils/
: Contains utility functions for reading/writing videos and analysing them.bbox_utils.py
helper.py
video.py
tracking/
: Contains implementation of object tracking algorithms.tracker.py
player_ball_assigner/
: Contains logic for assigning the ball to the closest player.player_ball_assigner.py
team_assigner/
: Contains algorithms for assigning players to teams based on jersey color.team_assigner.py
input/
: Directory for input video files.output/
: Directory for annotated output videos.models/
: Directory for storing trained models.stubs/
: Directory for storing pre-computed tracking data. This is where the pkl files generated for a video are stored. If a new input with same filename is given, instead of computing the tracking data again, it uses the pkl files. (Can be turned off by setting "read_from_stub=False" in main.py)development_and_analysis/
: Contains notebooks for testing and analysis.color_assignment.ipynb
training/
: Contains model training notebook to be used separately from the project to generate weights of the model.football_training.ipynb
-
Object Detection and Tracking: The
Tracker
class utilizes a pre-trained object detection model to detect players, referees, and the ball in each frame. It then tracks these objects across frames using tracking algorithms. -
Ball Interpolation: The ball might not be detected in some of the frames, hence, the missing ball positions are interpolated using the
interpolate_ball_positions
method to ensure smooth tracking throughout the video. -
Team Assignment: The
TeamAssigner
class segments player jersey colours using K-means clustering and assigns players to teams based on their jersey colours. -
Ball Assignment: The
PlayerBallAssigner
class determines the player closest to the ball and assigns ball possession accordingly. -
Annotation and Output: The processed frames are annotated with player IDs, teams, ball possession, and other relevant information. The annotated frames are then compiled into an output video file.
- YOLOv5 for object detection
- Abdullah Tarek for the project tutorial