init commit

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, exclude_objects = ['ball', 'referees']):
+        total_distance= {}
+
+        for object, object_tracks in tracks.items():
+            if object in exclude_objects:
+                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 == "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"{speed:.2f} km/h",position,cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,0),2)
+                       cv2.putText(frame, f"{distance:.2f} m",(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,73 @@
+from sklearn.cluster import KMeans
+
+class TeamAssigner:
+    def __init__(self):
+        self.team_colors = {}
+        self.player_team_dict = {}
+    
+    def get_clustering_model(self,image):
+        # Reshape the image to 2D array
+        image_2d = image.reshape(-1,3)
+
+        # Preform K-means with 2 clusters
+        kmeans = KMeans(n_clusters=2, init="k-means++",n_init=1)
+        kmeans.fit(image_2d)
+
+        return kmeans
+
+    def get_player_color(self,frame,bbox):
+        image = frame[int(bbox[1]):int(bbox[3]),int(bbox[0]):int(bbox[2])]
+
+        top_half_image = image[0:int(image.shape[0]/2),:]
+
+        # Get Clustering model
+        kmeans = self.get_clustering_model(top_half_image)
+
+        # Get the cluster labels forr each pixel
+        labels = kmeans.labels_
+
+        # Reshape the labels to the image shape
+        clustered_image = labels.reshape(top_half_image.shape[0],top_half_image.shape[1])
+
+        # Get the player cluster
+        corner_clusters = [clustered_image[0,0],clustered_image[0,-1],clustered_image[-1,0],clustered_image[-1,-1]]
+        non_player_cluster = max(set(corner_clusters),key=corner_clusters.count)
+        player_cluster = 1 - non_player_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)
+        kmeans.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
+
+        if player_id ==91:
+            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,218 @@
+from ultralytics import YOLO
+import supervision as sv
+import pickle
+import os
+import numpy as np
+import pandas as pd
+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_position_to_tracks(sekf,tracks):
+        for object, object_tracks in tracks.items():
+            for frame_num, track in enumerate(object_tracks):
+                for track_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][track_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 missing values
+        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=20 
+        detections = [] 
+        for i in range(0,len(frames),batch_size):
+            detections_batch = self.model.predict(frames[i:i+batch_size],conf=0.1)
+            detections += detections_batch
+        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":[]
+        }
+        print(len(tracks['players']), len(tracks['referees']), len(tracks['ball']))
+
+        for frame_num, detection in enumerate(detections):
+            cls_names = detection.names
+            cls_names_inv = {v:k for k,v in cls_names.items()}
+
+            # Covert to supervision Detection format
+            detection_supervision = sv.Detections.from_ultralytics(detection)
+
+            # Convert GoalKeeper to player object
+            for object_ind , class_id in enumerate(detection_supervision.class_id):
+                if cls_names[class_id] == "goalkeeper":
+                    detection_supervision.class_id[object_ind] = cls_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 == cls_names_inv['player']:
+                    tracks["players"][frame_num][track_id] = {"bbox":bbox}
+                
+                if cls_id == cls_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 == cls_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=(x_center,y2),
+            axes=(int(width), int(0.35*width)),
+            angle=0.0,
+            startAngle=-45,
+            endAngle=235,
+            color = color,
+            thickness=2,
+            lineType=cv2.LINE_4
+        )
+
+        rectangle_width = 40
+        rectangle_height=20
+        x1_rect = x_center - rectangle_width//2
+        x2_rect = x_center + rectangle_width//2
+        y1_rect = (y2- rectangle_height//2) +15
+        y2_rect = (y2+ rectangle_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_traingle(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 a semi-transparent rectaggle 
+        overlay = frame.copy()
+        cv2.rectangle(overlay, (1350, 850), (1900,970), (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 time each team had ball control
+        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 Possession: {team_1*100:.2f}%",(1400,900), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,0), 3)
+        cv2.putText(frame, f"Team 2 Possession: {team_2*100:.2f}%",(1400,950), cv2.FONT_HERSHEY_SIMPLEX, 1, (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]
+            ball_dict = tracks["ball"][frame_num]
+            referee_dict = tracks["referees"][frame_num]
+
+            # Draw Players
+            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_traingle(frame, player["bbox"],(0,0,255))
+
+            # Draw Referee
+            for _, referee in referee_dict.items():
+                frame = self.draw_ellipse(frame, referee["bbox"],(0,255,255))
+            
+            # Draw ball 
+            for track_id, ball in ball_dict.items():
+                frame = self.draw_traingle(frame, ball["bbox"],(255,255,50))
+
+
+            # 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 read_video, save_video
+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,18 @@
+import cv2
+
+def read_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        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,45 @@
+import numpy as np 
+import cv2
+
+class ViewTransformer():
+    def __init__(self):
+        court_width = 68
+        court_length = 23.32
+
+        self.pixel_vertices = np.array([[110, 1035], 
+                               [265, 275], 
+                               [910, 260], 
+                               [1640, 915]])
+        
+        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