File size: 10,764 Bytes
3d1f2c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 |
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):
# Generate all pairs of points in the segment
point_pairs = list(itertools.combinations(segment, 2))
# Calculate the Euclidean distance for each pair of points
distances = [math.dist((pair[0]['x'], pair[0]['y']), (pair[1]['x'], pair[1]['y'])) for pair in point_pairs]
# Find the indices of the pair with the maximum distance
max_distance_index = distances.index(max(distances))
# Extract the points from the pair with the maximum distance
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]
|