diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ad68280 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +files/output/* diff --git a/files/input/corner2.mp4 b/files/input/corner2.mp4 new file mode 100644 index 0000000..eff7b52 Binary files /dev/null and b/files/input/corner2.mp4 differ diff --git a/yolov5/detect.py b/yolov5/detect.py index 0161f96..99eb95c 100644 --- a/yolov5/detect.py +++ b/yolov5/detect.py @@ -19,9 +19,7 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized from deep_sort_pytorch.utils.parser import get_config from deep_sort_pytorch.deep_sort import DeepSort - -palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) - +import player def bbox_rel(*xyxy): """" Calculates the relative bounding box from absolute pixel values. """ @@ -36,12 +34,7 @@ def bbox_rel(*xyxy): return x_c, y_c, w, h -def compute_color_for_labels(label): - """ - Simple function that adds fixed color depending on the class - """ - color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette] - return tuple(color) +players = {} def draw_boxes(img, bbox, identities=None, offset=(0, 0)): @@ -53,14 +46,26 @@ def draw_boxes(img, bbox, identities=None, offset=(0, 0)): y2 += offset[1] # box text and bar id = int(identities[i]) if identities is not None else 0 - color = compute_color_for_labels(id) - label = '{}{:d}'.format("", id) - t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0] - cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) - cv2.rectangle( - img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1) - cv2.putText(img, label, (x1, y1 + - t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2) + if id in players.keys(): + current_player = players.get(id) + # only if checking colors automatically: + current_player.assignTeam(players) + label = current_player.team + else: + # check color manually + # team, color = player.check_color_manual2(left_clicks,img,x1,x2,y1,y2) + + # check color automatically + color = player.detectPlayerColor(img,x1,x2,y1,y2) + + current_player = player.Player(id,color=color,x=x2-(x2-x1),y=y2) + label = "?" + players[id] = current_player + + # label = current_player.team + plot_one_box(box, img, label=label, color=(int(current_player.color[0]), int(current_player.color[1]), int(current_player.color[2])), line_thickness=1) + #plot_one_box(box, img, label=label, color=current_player.color, line_thickness=1) + return img @@ -242,7 +247,7 @@ def detect(save_img=False): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') + parser.add_argument('--weights', nargs='+', type=str, default='yolov5l.pt', help='model.pt path(s)') parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold') diff --git a/yolov5/player.py b/yolov5/player.py new file mode 100644 index 0000000..db8d942 --- /dev/null +++ b/yolov5/player.py @@ -0,0 +1,227 @@ +import numpy as np +import cv2 +import torch +import colorsys +from sklearn.cluster import KMeans +from collections import Counter + +class Player: + + def __init__(self,id,color=None,team=None,x=None,y=None): + self.id = id + self.x = x + self.y = y + self.team = team + self.color = color + + def updatePosition(self, x, y): + self.x = x + self.y = y + + def assignTeam(self, players): + if self.team is None: + temp_list = [] + for key in players: + temp_list.append(players[key].color) + color_matrix = np.vstack((temp_list)) + clt = KMeans(n_clusters=3) + clt.fit(color_matrix) + n_pixels = len(clt.labels_) + counter = Counter(clt.labels_) + perc = {} + for i in counter: + perc[i] = np.round(counter[i]/n_pixels, 2) + perc = dict(sorted(perc.items())) + + main_colors = clt.cluster_centers_ + + max_value = max(perc, key=perc.get) + med_temp = list(sorted(perc.values()))[-2] + med_value = list(perc.keys())[list(perc.values()).index(med_temp)] + min_value = min(perc, key=perc.get) + + # hsv_player = cv2.cvtColor(np.uint8([[self.color]]), cv2.COLOR_BGR2HSV) + + # bgr_max = np.uint8([[main_colors[max_value]]]) + # hsv_max = cv2.cvtColor(bgr_max, cv2.COLOR_BGR2HSV) + + # bgr_med = np.uint8([[main_colors[med_value]]]) + # hsv_med = cv2.cvtColor(bgr_med, cv2.COLOR_BGR2HSV) + + # bgr_min = np.uint8([[main_colors[min_value]]]) + # hsv_min = cv2.cvtColor(bgr_min, cv2.COLOR_BGR2HSV) + + # adjust_array = np.array([10, 10, 40]) + + # lower_team1 = np.subtract(hsv_max, adjust_array) + # upper_team1 = np.add(hsv_max, adjust_array) + + # lower_team2 = np.subtract(hsv_med, adjust_array) + # upper_team2 = np.add(hsv_med, adjust_array) + + # lower_team3 = np.subtract(hsv_min, adjust_array) + # upper_team3 = np.add(hsv_min, adjust_array) + + # mask_team1 = cv2.inRange(hsv_player, lower_team1, upper_team1) + # mask_team2 = cv2.inRange(hsv_player, lower_team2, upper_team2) + # mask_team3 = cv2.inRange(hsv_player, lower_team3, upper_team3) + + # nonZero1 = cv2.countNonZero(mask_team1) + # nonZero2 = cv2.countNonZero(mask_team2) + # nonZero3 = cv2.countNonZero(mask_team3) + + # maxNonZero = max(nonZero1, nonZero2, nonZero3) + + # if maxNonZero == nonZero1: + # self.team = 1 + # self.color = main_colors[max_value] + # elif maxNonZero == nonZero2: + # self.team = 2 + # self.color = main_colors[med_value] + # else: + # self.team = 3 + # self.color = main_colors[min_value] + + + distances = np.sqrt(np.sum((main_colors-self.color)**2,axis=1)) + + index_of_smallest = np.where(distances==np.amin(distances)) + smallest_distance = main_colors[index_of_smallest] + if np.all(smallest_distance == main_colors[max_value]): + self.color = smallest_distance.flatten() + self.team = "Team_1" + # print(self.color, self.team) + elif np.all(smallest_distance == main_colors[med_value]): + self.color = smallest_distance.flatten() + self.team = "Team_2" + # print(self.color, self.team) + else: + self.color = self.color + self.team = "Other" + # print(self.color, self.team) + + # print(smallest_distance) + # print (main_colors[max_value], main_colors[med_value], main_colors[min_value]) + +def k_means(img): + clt = KMeans(n_clusters=4) + clt = clt.fit(img.reshape(-1, 3)) + n_pixels = len(clt.labels_) + counter = Counter(clt.labels_) + perc = {} + for i in counter: + perc[i] = np.round(counter[i]/n_pixels, 2) + perc = dict(sorted(perc.items())) + + return perc, clt.cluster_centers_ + + +def detectPlayerColor(img,x1,x2,y1,y2): + crop = img[y1:y2, x1:x2] + height, width, channels = crop.shape + qrt = crop[int(height/4):int(height/2), int(width/5):int(width/1.25)] + + perc, colors = k_means(qrt) + max_value = max(perc, key=perc.get) + return colors[max_value] + +def check_color_manual(img,x1,x2,y1,y2): + hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) + crop = hsv[y1:y2, x1:x2] + + height, width, channels = crop.shape + qrt = crop[int(height/4):int(height/2), int(width/5):int(width/1.25)] + + # rgb + team1 = np.array([37,69,234]) + team2 = np.array([255,217,215]) + team3 = np.array([0,0,0]) + + # hsv + lower_team1 = np.array([-5, 225, 215]) + upper_team1 = np.array([15, 245, 295]) + + lower_team2 = np.array([108, 33, 215]) + upper_team2 = np.array([128, 53, 295]) + + lower_team3 = np.array([144, 25, 11]) + upper_team3 = np.array([164, 45, 91]) + + mask_team1 = cv2.inRange(qrt, lower_team1, upper_team1) + mask_team2 = cv2.inRange(qrt, lower_team2, upper_team2) + mask_team3 = cv2.inRange(qrt, lower_team3, upper_team3) + + # out1 = cv2.bitwise_and(crop, crop, mask=mask_team1) + # out2 = cv2.bitwise_and(crop, crop, mask=mask_team2) + + nonZero1 = cv2.countNonZero(mask_team1) + nonZero2 = cv2.countNonZero(mask_team2) + nonZero3 = cv2.countNonZero(mask_team3) + + maxNonZero = max(nonZero1, nonZero2, nonZero3) + + if maxNonZero == nonZero1: + team = 1 + color = team1 + elif maxNonZero == nonZero2: + team = 2 + color = team2 + else: + team = 3 + color = team3 + + return (team, color) + +def check_color_manual2(clicks,img,x1,x2,y1,y2): + hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) + crop = hsv[y1:y2, x1:x2] + + height, width, channels = crop.shape + qrt = crop[int(height/4):int(height/2), int(width/5):int(width/1.25)] + + team1 = np.array(clicks[0]) + team2 = np.array(clicks[1]) + team3 = np.array(clicks[2]) + + bgr_A = np.uint8([[clicks[0]]]) + hsv_A = cv2.cvtColor(bgr_A, cv2.COLOR_BGR2HSV) + + bgr_B = np.uint8([[clicks[1]]]) + hsv_B = cv2.cvtColor(bgr_B, cv2.COLOR_BGR2HSV) + + bgr_C = np.uint8([[clicks[2]]]) + hsv_C = cv2.cvtColor(bgr_C, cv2.COLOR_BGR2HSV) + + adjust_array = np.array([10, 10, 40]) + + lower_team1 = np.subtract(hsv_A, adjust_array) + upper_team1 = np.add(hsv_A, adjust_array) + + lower_team2 = np.subtract(hsv_B, adjust_array) + upper_team2 = np.add(hsv_B, adjust_array) + + lower_team3 = np.subtract(hsv_C, adjust_array) + upper_team3 = np.add(hsv_C, adjust_array) + + mask_team1 = cv2.inRange(qrt, lower_team1, upper_team1) + mask_team2 = cv2.inRange(qrt, lower_team2, upper_team2) + mask_team3 = cv2.inRange(qrt, lower_team3, upper_team3) + + nonZero1 = cv2.countNonZero(mask_team1) + nonZero2 = cv2.countNonZero(mask_team2) + nonZero3 = cv2.countNonZero(mask_team3) + + maxNonZero = max(nonZero1, nonZero2, nonZero3) + + if maxNonZero == nonZero1: + team = 1 + color = team1 + elif maxNonZero == nonZero2: + team = 2 + color = team2 + else: + team = 3 + color = team3 + # print (color) + return (team, color) +