Spaces:
Sleeping
Sleeping
File size: 2,298 Bytes
7fa2fcb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
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
|