forked from Niketkumardheeryan/ML-CaPsule
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
adding remaining utilities for speed, movement, analyzing.
- Loading branch information
1 parent
f78014c
commit cb4a189
Showing
14 changed files
with
228 additions
and
0 deletions.
There are no files selected for viewing
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file not shown.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
Please assure training your model through the traning folder ipynb provided , by running it on colab and then storing the trained model in models folder. | ||
It is not possible to add the trained model here as it will overseed the github resource limit and is around 200mb in size. | ||
just follow the steps in training section and train your own model thorugh colab free quota easily. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
#upon running the model , the output video will appear here. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
from .player_ball_assigner import PlayerBallAssigner |
27 changes: 27 additions & 0 deletions
27
Football_Analyser_using_YOLO/player_ball_assigner/player_ball_assigner.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
1 change: 1 addition & 0 deletions
1
Football_Analyser_using_YOLO/speed_and_distance_estimator/__init__.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
from .speed_and_distance_estimator import SpeedAndDistance_Estimator |
73 changes: 73 additions & 0 deletions
73
Football_Analyser_using_YOLO/speed_and_distance_estimator/speed_and_distance_estimator.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
The training weights/stubs will appear here as a pkl file upon the first training run of model. | ||
Then for further runs those stubs will be uitlized saving the computational power. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
from .team_assigner import TeamAssigner |
73 changes: 73 additions & 0 deletions
73
Football_Analyser_using_YOLO/team_assigner/team_assigner.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
from .view_transformer import ViewTransformer |
45 changes: 45 additions & 0 deletions
45
Football_Analyser_using_YOLO/view_transformer/view_transformer.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |