|
|
import sys
|
|
|
import math
|
|
|
import itertools
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
from utils.utils_geometry import line_intersection
|
|
|
from utils.utils_heatmap import generate_gaussian_array_vectorized_l
|
|
|
|
|
|
|
|
|
class LineKeypointsDB(object):
|
|
|
def __init__(self, data, image):
|
|
|
|
|
|
self.lines_list = ["Big rect. left bottom",
|
|
|
"Big rect. left main",
|
|
|
"Big rect. left top",
|
|
|
"Big rect. right bottom",
|
|
|
"Big rect. right main",
|
|
|
"Big rect. right top",
|
|
|
"Goal left crossbar",
|
|
|
"Goal left post left ",
|
|
|
"Goal left post right",
|
|
|
"Goal right crossbar",
|
|
|
"Goal right post left",
|
|
|
"Goal right post right",
|
|
|
"Middle line",
|
|
|
"Side line bottom",
|
|
|
"Side line left",
|
|
|
"Side line right",
|
|
|
"Side line top",
|
|
|
"Small rect. left bottom",
|
|
|
"Small rect. left main",
|
|
|
"Small rect. left top",
|
|
|
"Small rect. right bottom",
|
|
|
"Small rect. right main",
|
|
|
"Small rect. right top"]
|
|
|
|
|
|
self.keypoint_pair_list = [['Side line top', 'Side line left'],
|
|
|
['Side line top', 'Middle line'],
|
|
|
['Side line right', 'Side line top'],
|
|
|
['Side line left', 'Big rect. left top'],
|
|
|
['Big rect. left top', 'Big rect. left main'],
|
|
|
['Big rect. right top', 'Big rect. right main'],
|
|
|
['Side line right', 'Big rect. right top'],
|
|
|
['Side line left', 'Small rect. left top'],
|
|
|
['Small rect. left top', 'Small rect. left main'],
|
|
|
['Small rect. right top', 'Small rect. right main'],
|
|
|
['Side line right', 'Small rect. right top'],
|
|
|
['Goal left crossbar', 'Goal left post right'],
|
|
|
['Side line left', 'Goal left post right'],
|
|
|
['Side line right', 'Goal right post left'],
|
|
|
['Goal right crossbar', 'Goal right post left'],
|
|
|
['Goal left crossbar', 'Goal left post left '],
|
|
|
['Side line left', 'Goal left post left '],
|
|
|
['Side line right', 'Goal right post right'],
|
|
|
['Goal right crossbar', 'Goal right post right'],
|
|
|
['Side line left', 'Small rect. left bottom'],
|
|
|
['Small rect. left bottom', 'Small rect. left main'],
|
|
|
['Small rect. right bottom', 'Small rect. right main'],
|
|
|
['Side line right', 'Small rect. right bottom'],
|
|
|
['Side line left', 'Big rect. left bottom'],
|
|
|
['Big rect. left bottom', 'Big rect. left main'],
|
|
|
['Big rect. right main', 'Big rect. right bottom'],
|
|
|
['Side line right', 'Big rect. right bottom'],
|
|
|
['Side line left', 'Side line bottom'],
|
|
|
['Side line bottom', 'Middle line'],
|
|
|
['Side line bottom', 'Side line right']]
|
|
|
|
|
|
self.line_keypoints_dict = {1: [24, 25],
|
|
|
2: [5, 25],
|
|
|
3: [4, 5],
|
|
|
4: [26, 27],
|
|
|
5: [6, 26],
|
|
|
6: [6, 7],
|
|
|
7: [12, 16],
|
|
|
8: [16, 17],
|
|
|
9: [12, 13],
|
|
|
10: [15, 19],
|
|
|
11: [14, 15],
|
|
|
12: [18, 19],
|
|
|
13: [2, 29],
|
|
|
14: [28, 30],
|
|
|
15: [1, 28],
|
|
|
16: [3, 30],
|
|
|
17: [1, 3],
|
|
|
18: [20, 21],
|
|
|
19: [9, 21],
|
|
|
20: [8, 9],
|
|
|
21: [22, 23],
|
|
|
22: [10, 22],
|
|
|
23: [10, 11]}
|
|
|
|
|
|
self.data = data
|
|
|
self.image = image
|
|
|
_, self.h, self.w = self.image.size()
|
|
|
self.size = (self.w, self.h)
|
|
|
|
|
|
self.num_channels = len(self.lines_list)
|
|
|
self.lines = {}
|
|
|
self.keypoints = {}
|
|
|
|
|
|
def get_tensor(self):
|
|
|
self.get_lines()
|
|
|
self.refine_point_lines()
|
|
|
heatmap_tensor = generate_gaussian_array_vectorized_l(self.num_channels, self.lines, self.size, down_ratio=2,
|
|
|
sigma=2)
|
|
|
return heatmap_tensor
|
|
|
|
|
|
def find_most_distanced_points(self, segment):
|
|
|
|
|
|
point_pairs = list(itertools.combinations(segment, 2))
|
|
|
|
|
|
|
|
|
distances = [math.dist((pair[0]['x'], pair[0]['y']), (pair[1]['x'], pair[1]['y'])) for pair in point_pairs]
|
|
|
|
|
|
|
|
|
max_distance_index = distances.index(max(distances))
|
|
|
|
|
|
|
|
|
most_distanced_points = list(point_pairs[max_distance_index])
|
|
|
|
|
|
return most_distanced_points
|
|
|
|
|
|
def get_main_keypoints(self):
|
|
|
for count, pair in enumerate(self.keypoint_pair_list):
|
|
|
if all(x in self.data.keys() for x in pair):
|
|
|
x, y = line_intersection(self.data, pair, self.w, self.h)
|
|
|
if not np.isnan(x):
|
|
|
if (0 <= x < self.w and 0 <= y < self.h):
|
|
|
self.keypoints[count + 1] = {'x': x, 'y': y, 'in_frame': True}
|
|
|
else:
|
|
|
self.keypoints[count + 1] = {'x': x, 'y': y, 'in_frame': False}
|
|
|
else:
|
|
|
self.keypoints[count + 1] = {'in_frame': False}
|
|
|
else:
|
|
|
self.keypoints[count + 1] = {'in_frame': False}
|
|
|
|
|
|
def get_lines(self):
|
|
|
for line in self.data.keys():
|
|
|
if line in self.lines_list:
|
|
|
if len(self.data[line]) > 1:
|
|
|
points = self.find_most_distanced_points(self.data[line])
|
|
|
x1, y1, x2, y2 = points[0]['x'], points[0]['y'], points[1]['x'], points[1]['y']
|
|
|
self.lines[self.lines_list.index(line) + 1] = {'x_1': x1 * self.w, 'y_1': y1 * self.h,
|
|
|
'x_2': x2 * self.w, 'y_2': y2 * self.h}
|
|
|
|
|
|
def refine_point_lines(self):
|
|
|
self.get_main_keypoints()
|
|
|
for line in self.lines.keys():
|
|
|
p1 = np.array([self.lines[line]['x_1'], self.lines[line]['y_1']])
|
|
|
p2 = np.array([self.lines[line]['x_2'], self.lines[line]['y_2']])
|
|
|
|
|
|
kp_to_refine = self.line_keypoints_dict[line]
|
|
|
kp1 = np.array([self.keypoints[kp_to_refine[0]]['x'], self.keypoints[kp_to_refine[0]]['y']]) if \
|
|
|
self.keypoints[kp_to_refine[0]]['in_frame'] else \
|
|
|
np.array([np.nan, np.nan])
|
|
|
kp2 = np.array([self.keypoints[kp_to_refine[1]]['x'], self.keypoints[kp_to_refine[1]]['y']]) if \
|
|
|
self.keypoints[kp_to_refine[1]]['in_frame'] else \
|
|
|
np.array([np.nan, np.nan])
|
|
|
|
|
|
|
|
|
if all(np.isnan(kp1)) and all(np.isnan(kp2)):
|
|
|
continue
|
|
|
|
|
|
elif not all(np.isnan(kp1)) and all(np.isnan(kp2)):
|
|
|
dist1 = np.linalg.norm(p1 - kp1)
|
|
|
dist2 = np.linalg.norm(p2 - kp1)
|
|
|
|
|
|
if dist1 < dist2:
|
|
|
self.lines[line]['x_1'] = kp1[0]
|
|
|
self.lines[line]['y_1'] = kp1[1]
|
|
|
else:
|
|
|
self.lines[line]['x_2'] = kp1[0]
|
|
|
self.lines[line]['y_2'] = kp1[1]
|
|
|
|
|
|
elif all(np.isnan(kp1)) and not all(np.isnan(kp2)):
|
|
|
dist1 = np.linalg.norm(p1 - kp2)
|
|
|
dist2 = np.linalg.norm(p2 - kp2)
|
|
|
|
|
|
if dist1 < dist2:
|
|
|
self.lines[line]['x_1'] = kp2[0]
|
|
|
self.lines[line]['y_1'] = kp2[1]
|
|
|
else:
|
|
|
self.lines[line]['x_2'] = kp2[0]
|
|
|
self.lines[line]['y_2'] = kp2[1]
|
|
|
else:
|
|
|
dis11 = np.linalg.norm(p1 - kp1)
|
|
|
dis12 = np.linalg.norm(p1 - kp2)
|
|
|
dis21 = np.linalg.norm(p2 - kp1)
|
|
|
dis22 = np.linalg.norm(p2 - kp2)
|
|
|
|
|
|
min_distance_p1 = min(dis11, dis12)
|
|
|
min_distance_p2 = min(dis21, dis22)
|
|
|
|
|
|
if min_distance_p1 < min_distance_p2:
|
|
|
if dis11 < dis12:
|
|
|
self.lines[line]['x_1'] = kp1[0]
|
|
|
self.lines[line]['y_1'] = kp1[1]
|
|
|
self.lines[line]['x_2'] = kp2[0]
|
|
|
self.lines[line]['y_2'] = kp2[1]
|
|
|
else:
|
|
|
self.lines[line]['x_1'] = kp2[0]
|
|
|
self.lines[line]['y_1'] = kp2[1]
|
|
|
self.lines[line]['x_2'] = kp1[0]
|
|
|
self.lines[line]['y_2'] = kp1[1]
|
|
|
else:
|
|
|
if dis11 < dis21:
|
|
|
self.lines[line]['x_1'] = kp1[0]
|
|
|
self.lines[line]['y_1'] = kp1[1]
|
|
|
self.lines[line]['x_2'] = kp2[0]
|
|
|
self.lines[line]['y_2'] = kp2[1]
|
|
|
else:
|
|
|
self.lines[line]['x_1'] = kp2[0]
|
|
|
self.lines[line]['y_1'] = kp2[1]
|
|
|
self.lines[line]['x_2'] = kp1[0]
|
|
|
self.lines[line]['y_2'] = kp1[1]
|
|
|
|
