| 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]
|
|
|
