color_detect

.gitignore

@@ -0,0 +1 @@
+files/output/*

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
 
-
+import player
-palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
-
 
 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):
+players = {}
-    """
-    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)
 
 
 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)
+        if id in players.keys():
-        label = '{}{:d}'.format("", id)
+            current_player = players.get(id)
-        t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
+            # only if checking colors automatically:
-        cv2.rectangle(img, (x1, y1), (x2, y2), color, 3)
+            current_player.assignTeam(players)
-        cv2.rectangle(
+            label = current_player.team
-            img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1)
+        else:
-        cv2.putText(img, label, (x1, y1 +
+            # check color manually
-                                 t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2)
+            # 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')

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)
+