commitSecond

adding utility for camera movement, color assignment, ball assignment , and speed and distance assignment.

FootballAnalysisWithYolo/camera_movement_estimator/__init__.py

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

FootballAnalysisWithYolo/camera_movement_estimator/camera_movement_estimator.py

@@ -0,0 +1,99 @@
+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[:,900:1050] = 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 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):
+            frame= frame.copy()
+
+            overlay = frame.copy()
+            cv2.rectangle(overlay,(0,0),(500,100),(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,30), cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,0),3)
+            frame = cv2.putText(frame,f"Camera Movement Y: {y_movement:.2f}",(10,60), cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,0),3)
+
+            output_frames.append(frame) 
+
+        return output_frames

FootballAnalysisWithYolo/player_ball_assigner/__init__.py

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

FootballAnalysisWithYolo/player_ball_assigner/player_ball_assigner.py

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

FootballAnalysisWithYolo/speed_and_distance_estimator/__init__.py

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

FootballAnalysisWithYolo/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"{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