s406917/VisionScore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

color_detect

Browse Source
This commit is contained in:
bednarco 2021-01-17 20:55:54 +01:00
parent f3b0d78b65
commit 27136567d4
4 changed files with 251 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored Normal file
View File

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

BIN
files/input/corner2.mp4 Normal file
View File

Binary file not shown.

41
yolov5/detect.py
View File

@ -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')

227
yolov5/player.py Normal file
View File

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