let's deploy to huggingface spaces

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,167 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+
+.vscode/
+runs/
+development_and_analysis/
+stubs/
+training/
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

README.md

@@ -1,6 +1,6 @@
 ---
 title: Football Analysis
-emoji: 🏢
+emoji: ⚽
 colorFrom: purple
 colorTo: red
 sdk: gradio
@@ -11,3 +11,11 @@ license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Similar Input Videos
+
+<https://www.kaggle.com/competitions/dfl-bundesliga-data-shootout/data?select=clips>
+
+The aim of this project is to detect and track players, referees, and footballs in video footage using YOLO, a top-tier AI object detection model. ⚽ We will enhance the model's performance through training. Additionally, we will classify players into teams based on their t-shirt colors using K-means for pixel segmentation and clustering. 🎽📊 This data will enable us to calculate a team's ball possession percentage during a match.
+
+To accurately measure player movement, we will use optical flow to assess camera motion between frames and apply perspective transformation to represent the scene's depth and perspective. This allows us to convert player movement from pixels to meters. 📐 Finally, we will calculate each player's distance covered. This comprehensive project encompasses a range of concepts and tackles real-world problems, making it ideal for both novice and experienced machine learning engineers. 🚀

app.py

@@ -0,0 +1,102 @@
+from click import launch
+from player_ball_assigner import PlayerBallAssigner
+from utils import read_video,save_video,check_video_resolution
+from trackers import Tracker
+from team_assigner import TeamAssigner
+import numpy as np
+from camera_movement_estimator import CameraMovementEstimator
+from view_transformer import ViewTransformer
+from speed_and_distance_estimator import SpeedAndDistance_Estimator
+import gradio as gr
+
+def annotate(input):
+ 
+    # Read Video
+    video_frames = read_video(input)
+
+
+    #intialize tracker
+    tracker = Tracker('models/best.pt')
+    tracks = tracker.get_object_tracks(video_frames,read_from_stub=False,stub_path=None)
+
+    # Get Object Postions
+    tracker.add_postion_to_tracks(tracks)
+
+    # camera movement estimator
+    camera_movement_estimator = CameraMovementEstimator(video_frames[0])
+    camera_movement_per_frame = camera_movement_estimator.get_camera_movement(video_frames,
+                                                                                read_from_stub=False,
+                                                                                stub_path=None)
+    camera_movement_estimator.add_adjust_positions_to_tracks(tracks,camera_movement_per_frame)
+
+        # View Trasnformer
+    view_transformer = ViewTransformer()
+    view_transformer.add_transformed_position_to_tracks(tracks)
+
+    # interpolate ball positions
+    tracks['ball'] = tracker.interpolate_ball_positions(tracks['ball'])
+
+
+    # Speed and distance estimator
+    speed_and_distance_estimator = SpeedAndDistance_Estimator()
+    speed_and_distance_estimator.add_speed_and_distance_to_tracks(tracks)
+
+    # Assign player teams
+    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]
+            
+
+
+    # Assign Ball Acquisition
+    player_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_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)
+
+    # Draw Output
+    ## Draw object tracks
+    output_video_frames = tracker.draw_annotations(video_frames,tracks,team_ball_control)
+
+     ## Draw Camera movement
+    output_video_frames = camera_movement_estimator.draw_camera_movement(output_video_frames,camera_movement_per_frame)
+
+    ## Draw Speed and Distance
+    speed_and_distance_estimator.draw_speed_and_distance(output_video_frames,tracks)
+
+    output_path= 'output_videos/out.mp4'
+    # Save Video
+    save_video(output_video_frames,output_path)
+
+
+    return output_path
+
+
+
+
+iface = gr.Interface(
+    fn=annotate,
+    inputs=['video'],
+    outputs=['video'],
+    title="Football Analysis",
+    examples=['input_videos/08fd33_4.mp4'],
+
+).queue(default_concurrency_limit=2).launch()
+

camera_movement_estimator/__init__.py

@@ -0,0 +1 @@
+from .camera_movement_estimator import CameraMovementEstimator

camera_movement_estimator/camera_movement_estimator.py

@@ -0,0 +1,108 @@
+import pickle
+import cv2
+import numpy as np
+import os
+import sys 
+sys.path.append('../')
+from utils import measure_distance,measure_xy_distance
+
+class CameraMovementEstimator():
+    def __init__(self,frame):
+        self.minimum_distance = 5
+
+        self.lk_params = dict(
+            winSize = (15,15),
+            maxLevel = 2,
+            criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,10,0.03)
+        )
+
+        first_frame_grayscale = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
+        mask_features = np.zeros_like(first_frame_grayscale)
+        mask_features[:,0:20] = 1
+        mask_features[:,600:700] = 1
+
+        self.features = dict(
+            maxCorners = 100,
+            qualityLevel = 0.3,
+            minDistance =3,
+            blockSize = 7,
+            mask = mask_features
+        )
+
+    
+    def add_adjust_positions_to_tracks(self,tracks, camera_movement_per_frame):
+        for object, object_tracks in tracks.items():
+            for frame_num, track in enumerate(object_tracks):
+                for track_id, track_info in track.items():
+                    position = track_info['position']
+                    camera_movement = camera_movement_per_frame[frame_num]
+                    position_adjusted = (position[0]-camera_movement[0],position[1]-camera_movement[1])
+                    tracks[object][frame_num][track_id]['position_adjusted'] = position_adjusted
+
+    def add_adjust_positions_to_tracks(self,tracks, camera_movement_per_frame):
+        for object, object_tracks in tracks.items():
+            for frame_num, track in enumerate(object_tracks):
+                for track_id, track_info in track.items():
+                    position = track_info['position']
+                    camera_movement = camera_movement_per_frame[frame_num]
+                    position_adjusted = (position[0]-camera_movement[0],position[1]-camera_movement[1])
+                    tracks[object][frame_num][track_id]['position_adjusted'] = position_adjusted
+                    
+
+
+    def get_camera_movement(self,frames,read_from_stub=False, stub_path=None):
+        # Read the stub 
+        if read_from_stub and stub_path is not None and os.path.exists(stub_path):
+            with open(stub_path,'rb') as f:
+                return pickle.load(f)
+
+        camera_movement = [[0,0]]*len(frames)
+
+        old_gray = cv2.cvtColor(frames[0],cv2.COLOR_BGR2GRAY)
+        old_features = cv2.goodFeaturesToTrack(old_gray,**self.features)
+
+        for frame_num in range(1,len(frames)):
+            frame_gray = cv2.cvtColor(frames[frame_num],cv2.COLOR_BGR2GRAY)
+            new_features, _,_ = cv2.calcOpticalFlowPyrLK(old_gray,frame_gray,old_features,None,**self.lk_params)
+
+            max_distance = 0
+            camera_movement_x, camera_movement_y = 0,0
+
+            for i, (new,old) in enumerate(zip(new_features,old_features)):
+                new_features_point = new.ravel()
+                old_features_point = old.ravel()
+
+                distance = measure_distance(new_features_point,old_features_point)
+                if distance>max_distance:
+                    max_distance = distance
+                    camera_movement_x,camera_movement_y = measure_xy_distance(old_features_point, new_features_point ) 
+            
+            if max_distance > self.minimum_distance:
+                camera_movement[frame_num] = [camera_movement_x,camera_movement_y]
+                old_features = cv2.goodFeaturesToTrack(frame_gray,**self.features)
+
+            old_gray = frame_gray.copy()
+        
+        if stub_path is not None:
+            with open(stub_path,'wb') as f:
+                pickle.dump(camera_movement,f)
+
+        return camera_movement
+    
+    def draw_camera_movement(self,frames, camera_movement_per_frame):
+        output_frames=[]
+
+        for frame_num, frame in enumerate(frames):
+
+            overlay = frame.copy()
+            cv2.rectangle(overlay,(0,0),(350,70),(255,255,255),-1)
+            alpha =0.6
+            cv2.addWeighted(overlay,alpha,frame,1-alpha,0,frame)
+
+            x_movement, y_movement = camera_movement_per_frame[frame_num]
+            frame = cv2.putText(frame,f"Camera Movement X: {x_movement:.2f}",(10,20), cv2.FONT_HERSHEY_SIMPLEX,0.7,(0,0,0),3)
+            frame = cv2.putText(frame,f"Camera Movement Y: {y_movement:.2f}",(10,60), cv2.FONT_HERSHEY_SIMPLEX,0.7,(0,0,0),3)
+
+            output_frames.append(frame) 
+
+        return output_frames

input_videos/08fd33_4.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:077985a55b61fe2ca413221c082711655b7ba736a3bbd8ff97b4dfc0e5b03c57
+size 20341580

models/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ffeb0ff661427e8cbe0763ebfba72d385f18737c5e0070d36e911a586f5e7bc
+size 194957941

player_ball_assigner/__init__.py

@@ -0,0 +1 @@
+from .player_ball_assigner import PlayerBallAssigner

player_ball_assigner/player_ball_assigner.py

@@ -0,0 +1,26 @@
+import sys
+sys.path.append('../')
+from utils import get_center_of_bbox,measure_distance
+
+class PlayerBallAssigner():
+    def __init__(self) -> None:
+        self.max_player_ball_dist = 70
+
+    def assign_ball_to_player(self,players,ball_bbox):
+        ball_position = get_center_of_bbox(ball_bbox)
+
+        minimum_dist = 99999
+        assigned_player = -1
+
+        for player_id,player in players.items():
+            player_bbox = player['bbox']
+
+            dist_left = measure_distance((player_bbox[0],player_bbox[-1]),ball_position)
+            dist_right = measure_distance((player_bbox[2],player_bbox[-1]),ball_position)
+            distance = min(dist_left,dist_right)
+
+            if distance < self.max_player_ball_dist:
+                minimum_dist=distance
+                assigned_player=player_id
+            
+        return assigned_player

requirements.txt

@@ -0,0 +1,8 @@
+ultralytics
+cv2
+gradio
+numpy
+pandas
+pickle
+supervision
+sklearn

speed_and_distance_estimator/__init__.py

@@ -0,0 +1 @@
+from .speed_and_distance_estimator import SpeedAndDistance_Estimator

speed_and_distance_estimator/speed_and_distance_estimator.py

@@ -0,0 +1,73 @@
+import cv2
+import sys 
+sys.path.append('../')
+from utils import measure_distance ,get_foot_position
+
+class SpeedAndDistance_Estimator():
+    def __init__(self):
+        self.frame_window=5
+        self.frame_rate=24
+    
+    def add_speed_and_distance_to_tracks(self,tracks):
+        total_distance= {}
+
+        for object, object_tracks in tracks.items():
+            if object == "ball" or object == "referees":
+                continue 
+            number_of_frames = len(object_tracks)
+            for frame_num in range(0,number_of_frames, self.frame_window):
+                last_frame = min(frame_num+self.frame_window,number_of_frames-1 )
+
+                for track_id,_ in object_tracks[frame_num].items():
+                    if track_id not in object_tracks[last_frame]:
+                        continue
+
+                    start_position = object_tracks[frame_num][track_id]['position_transformed']
+                    end_position = object_tracks[last_frame][track_id]['position_transformed']
+
+                    if start_position is None or end_position is None:
+                        continue
+                    
+                    distance_covered = measure_distance(start_position,end_position)
+                    time_elapsed = (last_frame-frame_num)/self.frame_rate
+                    speed_meteres_per_second = distance_covered/time_elapsed
+                    speed_km_per_hour = speed_meteres_per_second*3.6
+
+                    if object not in total_distance:
+                        total_distance[object]= {}
+                    
+                    if track_id not in total_distance[object]:
+                        total_distance[object][track_id] = 0
+                    
+                    total_distance[object][track_id] += distance_covered
+
+                    for frame_num_batch in range(frame_num,last_frame):
+                        if track_id not in tracks[object][frame_num_batch]:
+                            continue
+                        tracks[object][frame_num_batch][track_id]['speed'] = speed_km_per_hour
+                        tracks[object][frame_num_batch][track_id]['distance'] = total_distance[object][track_id]
+    
+    def draw_speed_and_distance(self,frames,tracks):
+        output_frames = []
+        for frame_num, frame in enumerate(frames):
+            for object, object_tracks in tracks.items():
+                if object == "ball" or object == "referees":
+                    continue 
+                for _, track_info in object_tracks[frame_num].items():
+                   if "speed" in track_info:
+                       speed = track_info.get('speed',None)
+                       distance = track_info.get('distance',None)
+                       if speed is None or distance is None:
+                           continue
+                       
+                       bbox = track_info['bbox']
+                       position = get_foot_position(bbox)
+                       position = list(position)
+                       position[1]+=40
+
+                       position = tuple(map(int,position))
+                       cv2.putText(frame, f"{distance:.2f} m",position,cv2.FONT_HERSHEY_SIMPLEX,0.4,(0,0,0),2)
+                    #    cv2.putText(frame, f"{speed:.2f} km/h",(position[0],position[1]+20),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,0),2)
+            output_frames.append(frame)
+        
+        return output_frames

team_assigner/__init__.py

@@ -0,0 +1 @@
+from .team_assigner import TeamAssigner

team_assigner/team_assigner.py

@@ -0,0 +1,72 @@
+from sklearn.cluster import KMeans
+
+class TeamAssigner:
+    def __init__(self) -> None:
+        self.team_colors = {}
+        self.player_team_dict = {}
+
+
+    def get_clustering_model(self,img):
+        # Reshape the image to 2D Array
+        img_2d = img.reshape(-1,3)
+
+        # Perform K-means clustering with 2 clusters
+        kmeans = KMeans(n_clusters=2,init="k-means++",n_init=1).fit(img_2d)
+
+        return kmeans
+
+
+
+    def get_player_color(self,frame,bbox):
+        img = frame[int(bbox[1]):int(bbox[3]),int(bbox[0]):int(bbox[2])]
+
+        top_img_hlf = img[0:int(img.shape[0]/2),:]
+
+        # Get Clustering model
+        kmeans = self.get_clustering_model(top_img_hlf)    
+
+        # Get the cluster labels for each pixel
+        labels = kmeans.labels_
+
+        # Reshape the labels to original img shape
+        clustered_img = labels.reshape(top_img_hlf.shape[0],top_img_hlf.shape[1])
+
+        # Get the player cluster
+        corner_clusters = [clustered_img[0,0],clustered_img[0,-1],clustered_img[-1,0],clustered_img[-1,-1]]
+        np_cluster = max(set(corner_clusters),key=corner_clusters.count)
+        player_cluster = 1-np_cluster
+
+        player_color=kmeans.cluster_centers_[player_cluster]
+
+        return player_color
+
+
+    def assign_team_color(self,frame,player_detections):
+        
+        player_colors = []
+        for _,player_detection in player_detections.items():
+            bbox = player_detection['bbox']
+            player_color = self.get_player_color(frame,bbox)
+            player_colors.append(player_color)
+
+        kmeans = KMeans(n_clusters=2,init='k-means++',n_init=10).fit(player_colors)
+        
+        self.kmeans = kmeans
+
+        self.team_colors[1] = kmeans.cluster_centers_[0]
+        self.team_colors[2] = kmeans.cluster_centers_[1]
+
+    
+    def get_player_team(self,frame,player_bbox,player_id):
+        if player_id in self.player_team_dict:
+            return self.player_team_dict[player_id]
+        
+        player_color = self.get_player_color(frame,player_bbox)
+
+        team_id = self.kmeans.predict(player_color.reshape(1,-1))[0]
+        team_id +=1
+
+        self.player_team_dict[player_id] = team_id
+
+        return team_id
+        

trackers/__init__.py

@@ -0,0 +1 @@
+from .tracker import Tracker

trackers/tracker.py

@@ -0,0 +1,241 @@
+from ultralytics import YOLO
+import supervision as sv
+import pickle
+import numpy as np
+import pandas as pd
+import os
+import cv2
+import sys
+
+sys.path.append('../')
+from utils import get_center_of_bbox,get_bbox_width,get_foot_position
+
+class Tracker:
+    def __init__(self,model_path):
+        self.model = YOLO(model_path)
+        self.tracker = sv.ByteTrack()
+
+    
+    def add_postion_to_tracks(self,tracks):
+        for object,object_tracks in tracks.items():
+            for frame_num,track in enumerate(object_tracks):
+                for tracks_id,track_info in track.items():
+                    bbox = track_info['bbox']
+                    if object == 'ball':
+                        position = get_center_of_bbox(bbox)
+                    else:
+                        position = get_foot_position(bbox)
+                    tracks[object][frame_num][tracks_id]['position'] = position
+    
+    
+
+
+    def interpolate_ball_positions(self,ball_positions):
+        ball_positions = [x.get(1,{}).get('bbox',[]) for x in ball_positions]
+        df_ball_positions = pd.DataFrame(ball_positions,columns=['x1','y1','x2','y2'])
+
+        # Interpolate ball postions 
+        df_ball_positions = df_ball_positions.interpolate()
+        df_ball_positions = df_ball_positions.bfill()
+
+        ball_positions = [{1:{"bbox":x}} for x in df_ball_positions.to_numpy().tolist()]
+
+        return ball_positions
+
+    def detect_frames(self,frames):
+        batch_size = 50
+        detections = []
+        for i in range(0,len(frames),batch_size):
+            detections_btach = self.model.predict(frames[i:i+batch_size],conf=0.1)
+            detections += detections_btach
+            
+        
+        return detections
+    
+    def get_object_tracks(self,frames,read_from_stub=False,stub_path=None):
+
+        if read_from_stub and stub_path is not None and os.path.exists(stub_path):
+            with open(stub_path,'rb') as f:
+                tracks = pickle.load(f)
+
+            return tracks
+        
+        detections = self.detect_frames(frames)
+
+        tracks = {
+            "players":[],
+            "referees":[],
+            "ball":[]
+        }
+
+        for frame_num , detection in enumerate(detections):
+            class_names = detection.names
+            class_names_inv = {v:k for k,v in class_names.items()}
+
+            # Convert to supervision detection format
+            detection_supervision = sv.Detections.from_ultralytics(detection)
+
+            # Convert goalkeeper to player object
+            for index , class_id in enumerate(detection_supervision.class_id):
+                if class_names[class_id] == 'goalkeeper':
+                    detection_supervision.class_id[index] = class_names_inv['player']
+
+            # Track Objects
+            detection_with_tracks = self.tracker.update_with_detections(detection_supervision)
+
+            tracks["players"].append({})
+            tracks["referees"].append({})
+            tracks['ball'].append({})
+
+            for frame_detection in detection_with_tracks:
+                bbox = frame_detection[0].tolist()
+                cls_id = frame_detection[3]
+                track_id = frame_detection[4]
+
+                if cls_id == class_names_inv['player']:
+                    tracks['players'][frame_num][track_id] = {"bbox":bbox}
+
+                if cls_id == class_names_inv['referee']:
+                    tracks['referees'][frame_num][track_id] = {"bbox":bbox}
+
+            for frame_detection in detection_supervision:
+                bbox = frame_detection[0].tolist()
+                cls_id = frame_detection[3]
+
+                if cls_id == class_names_inv['ball']:
+                    tracks['ball'][frame_num][1] = {"bbox":bbox}
+        
+        
+        if stub_path is not None:
+            with open(stub_path,'wb') as f:
+                pickle.dump(tracks,f)
+
+        return tracks
+    
+
+    def draw_ellipse(self,frame,bbox,color,track_id=None):
+        y2 = int(bbox[3])
+        x_center,_ = get_center_of_bbox(bbox)
+        width = get_bbox_width(bbox)
+
+        cv2.ellipse(
+            frame,
+            center=(int(x_center),int(y2)),
+            axes=(int(width),int(0.35*width)),
+            angle=0.0,
+            startAngle=-45.0,
+            endAngle=235.0,
+            color=color,
+            thickness=2,
+            lineType=cv2.LINE_4,
+        )
+
+        rect_width = 40
+        rect_height = 20
+        x1_rect = x_center - rect_width//2
+        x2_rect = x_center + rect_width//2
+        y1_rect = (y2 - rect_height//2) + 15
+        y2_rect = (y2 + rect_height//2) + 15
+
+        if track_id is not None:
+            cv2.rectangle(
+                frame,
+                (int(x1_rect),int(y1_rect)),
+                (int(x2_rect),int(y2_rect)),
+                color,
+                cv2.FILLED
+            )
+
+            x1_text = x1_rect + 12
+            if track_id > 99:
+                x1_text -=10
+
+            cv2.putText(
+                frame,
+                f"{track_id}",
+                (int(x1_text),int(y1_rect+15)),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.6,
+                (0,0,0),
+                2
+            )
+        
+
+        return frame
+    
+
+    def draw_triangle(self,frame,bbox,color):
+        y = int(bbox[1])
+        x,_ = get_center_of_bbox(bbox)
+
+        triangle_points = np.array([
+            [x,y],
+            [x-10,y-20],
+            [x+10,y-20],
+        ])
+
+        cv2.drawContours(frame,[triangle_points],0,color,cv2.FILLED)
+        cv2.drawContours(frame,[triangle_points],0,(0,0,0),2)
+
+        return frame
+    
+    def draw_team_ball_control(self,frame,frame_num,team_ball_control):
+        # Draw semi transparent rectangle
+        overlay = frame.copy()
+        cv2.rectangle(overlay,(900,570),(1250,660),(255,255,255),-1)
+        alpha=0.4
+        cv2.addWeighted(overlay,alpha,frame,1-alpha,0,frame)
+
+        team_ball_control_till_frame = team_ball_control[:frame_num+1]
+
+        # get the number of times each team had the ball
+        team_1_num_frames = team_ball_control_till_frame[team_ball_control_till_frame==1].shape[0]
+        team_2_num_frames = team_ball_control_till_frame[team_ball_control_till_frame==2].shape[0]
+        team_1 = team_1_num_frames/(team_1_num_frames+team_2_num_frames)
+        team_2 = team_2_num_frames/(team_1_num_frames+team_2_num_frames)
+
+        cv2.putText(frame,f"Team 1 Ball Control {team_1*100:.2f}%",(920,600),cv2.FONT_HERSHEY_SIMPLEX,0.7,(0,0,0),3)
+        cv2.putText(frame,f"Team 2 Ball Control {team_2*100:.2f}%",(920,640),cv2.FONT_HERSHEY_SIMPLEX,0.7,(0,0,0),3)
+
+        return frame
+
+
+    def draw_annotations(self,video_frames,tracks,team_ball_control):
+        output_video_frames = []
+        for frame_num,frame in enumerate(video_frames):
+            # frame = frame.copy()
+            
+            player_dict = tracks['players'][frame_num]
+            referee_dict = tracks['referees'][frame_num]
+            ball_dict = tracks['ball'][frame_num]
+
+            # Draw player annotations
+            for track_id,player in player_dict.items():
+                color = player.get('team_color',(0,0,255))
+                frame = self.draw_ellipse(frame,player['bbox'],color,track_id)
+
+                if player.get('has_ball',False):
+                    frame = self.draw_triangle(frame,player['bbox'],(0,0,255))
+
+
+            # Draw referee annotations
+            for _,referee in referee_dict.items():
+                frame = self.draw_ellipse(frame,referee['bbox'],(0,255,255))
+
+            # Draw ball annotaion
+
+            for _,ball in ball_dict.items():
+                frame = self.draw_triangle(frame,ball['bbox'],(0,255,0))
+
+            # Draw team ball control
+
+            frame = self.draw_team_ball_control(frame,frame_num,team_ball_control)
+
+            output_video_frames.append(frame)
+        
+        return output_video_frames
+            
+
+            
+            
+

utils/__init__.py

@@ -0,0 +1,2 @@
+from .video_utils import save_video,read_video,check_video_resolution
+from .bbox_utils import get_center_of_bbox,get_bbox_width,measure_distance,measure_xy_distance,get_foot_position

utils/bbox_utils.py

@@ -0,0 +1,16 @@
+def get_center_of_bbox(bbox):
+    x1,y1,x2,y2 = bbox 
+    return int((x1+x2)/2),int((y1+y2)/2)
+
+def get_bbox_width(bbox):
+    return bbox[2]-bbox[0]
+
+def measure_distance(p1,p2):
+    return ((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2)**0.5
+
+def measure_xy_distance(p1,p2):
+    return p1[0]-p2[0],p1[1]-p2[1]
+
+def get_foot_position(bbox):
+    x1,y1,x2,y2 = bbox
+    return int((x1+x2)/2),int(y2)

utils/video_utils.py

@@ -0,0 +1,26 @@
+import cv2
+
+
+def check_video_resolution(video_path):
+    vid = cv2.VideoCapture(video_path)
+    height = vid.get(cv2.CAP_PROP_FRAME_HEIGHT)
+    width = vid.get(cv2.CAP_PROP_FRAME_WIDTH)
+    return height,width
+
+def read_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.resize(frame,(1280,720))
+        frames.append(frame)
+    return frames
+
+def save_video(ouput_video_frames,output_video_path):
+    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    out = cv2.VideoWriter(output_video_path, fourcc, 24, (ouput_video_frames[0].shape[1], ouput_video_frames[0].shape[0]))
+    for frame in ouput_video_frames:
+        out.write(frame)
+    out.release()

view_transformer/__init__.py

@@ -0,0 +1 @@
+from .view_transformer import ViewTransformer

view_transformer/view_transformer.py

@@ -0,0 +1,46 @@
+import numpy as np 
+import cv2
+
+class ViewTransformer():
+    def __init__(self):
+        court_width = 68
+        court_length = 29.15
+
+        self.pixel_vertices = np.array([
+                               [78,689], 
+                               [177,182], 
+                               [699,169], 
+                               [1280,599]])
+        
+        self.target_vertices = np.array([
+            [0,court_width],
+            [0, 0],
+            [court_length, 0],
+            [court_length, court_width]
+        ])
+
+        self.pixel_vertices = self.pixel_vertices.astype(np.float32)
+        self.target_vertices = self.target_vertices.astype(np.float32)
+
+        self.persepctive_trasnformer = cv2.getPerspectiveTransform(self.pixel_vertices, self.target_vertices)
+
+    def transform_point(self,point):
+        p = (int(point[0]),int(point[1]))
+        is_inside = cv2.pointPolygonTest(self.pixel_vertices,p,False) >= 0 
+        if not is_inside:
+            return None
+
+        reshaped_point = point.reshape(-1,1,2).astype(np.float32)
+        tranform_point = cv2.perspectiveTransform(reshaped_point,self.persepctive_trasnformer)
+        return tranform_point.reshape(-1,2)
+
+    def add_transformed_position_to_tracks(self,tracks):
+        for object, object_tracks in tracks.items():
+            for frame_num, track in enumerate(object_tracks):
+                for track_id, track_info in track.items():
+                    position = track_info['position_adjusted']
+                    position = np.array(position)
+                    position_trasnformed = self.transform_point(position)
+                    if position_trasnformed is not None:
+                        position_trasnformed = position_trasnformed.squeeze().tolist()
+                    tracks[object][frame_num][track_id]['position_transformed'] = position_trasnformed