|
|
import sys
|
|
|
import cv2
|
|
|
import math
|
|
|
import copy
|
|
|
import torch
|
|
|
import itertools
|
|
|
import numpy as np
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
from mpl_toolkits.axes_grid1 import make_axes_locatable
|
|
|
from scipy.optimize import linear_sum_assignment
|
|
|
from scipy.stats import linregress
|
|
|
from ellipse import LsqEllipse
|
|
|
from itertools import product
|
|
|
from functools import reduce
|
|
|
|
|
|
from utils.utils_field import _draw_field
|
|
|
from utils.utils_heatmap import generate_gaussian_array_vectorized
|
|
|
|
|
|
|
|
|
class KeypointsWCDB(object):
|
|
|
def __init__(self, image, homography, size_out=(960,540)):
|
|
|
|
|
|
self.keypoint_world_coords_2D = [[0., 0.], [52.5, 0.], [105., 0.], [0., 13.84], [16.5, 13.84], [88.5, 13.84],
|
|
|
[105., 13.84], [0., 24.84], [5.5, 24.84], [99.5, 24.84], [105., 24.84],
|
|
|
[0., 30.34], [0., 30.34], [105., 30.34], [105., 30.34], [0., 37.66],
|
|
|
[0., 37.66], [105., 37.66], [105., 37.66], [0., 43.16], [5.5, 43.16],
|
|
|
[99.5, 43.16], [105., 43.16], [0., 54.16], [16.5, 54.16], [88.5, 54.16],
|
|
|
[105., 54.16], [0., 68.], [52.5, 68.], [105., 68.], [16.5, 26.68],
|
|
|
[52.5, 24.85], [88.5, 26.68], [16.5, 41.31], [52.5, 43.15], [88.5, 41.31],
|
|
|
[19.99, 32.29], [43.68, 31.53], [61.31, 31.53], [85., 32.29], [19.99, 35.7],
|
|
|
[43.68, 36.46], [61.31, 36.46], [85., 35.7], [11., 34.], [16.5, 34.],
|
|
|
[20.15, 34.], [46.03, 27.53], [58.97, 27.53], [43.35, 34.], [52.5, 34.],
|
|
|
[61.5, 34.], [46.03, 40.47], [58.97, 40.47], [84.85, 34.], [88.5, 34.],
|
|
|
[94., 34.]]
|
|
|
|
|
|
self.keypoint_aux_world_coords_2D = [[5.5, 0], [16.5, 0], [88.5, 0], [99.5, 0], [5.5, 13.84], [99.5, 13.84],
|
|
|
[16.5, 24.84], [88.5, 24.84], [16.5, 43.16], [88.5, 43.16], [5.5, 54.16],
|
|
|
[99.5, 54.16], [5.5, 68], [16.5, 68], [88.5, 68], [99.5, 68]]
|
|
|
|
|
|
self.lines_retrieval = [[24, 25], [5, 25], [4, 5], [26, 27], [6, 26], [12, 16], [16, 17], [12, 13], [15, 19],
|
|
|
[14, 15], [18, 19], [2, 29], [28, 29, 30], [1, 4, 8, 13, 17, 20, 24, 28],
|
|
|
[3, 7, 11, 14, 18, 23, 27, 30], [1, 2, 3], [20, 21], [9, 21], [8, 9], [22, 23],
|
|
|
[10, 22], [10, 11]]
|
|
|
|
|
|
|
|
|
self.homography = homography
|
|
|
self.image = image
|
|
|
|
|
|
self.w, self.h = size_out
|
|
|
self.size = (self.w, self.h)
|
|
|
self.h_extra = self.h * 0.5
|
|
|
self.w_extra = self.w * 0.5
|
|
|
|
|
|
self.keypoints_final = {}
|
|
|
|
|
|
self.num_channels = len(self.keypoint_world_coords_2D) + 1
|
|
|
self.mask_array = np.ones(self.num_channels).astype(int)
|
|
|
|
|
|
|
|
|
def get_tensor_w_mask(self):
|
|
|
|
|
|
self.get_kp_from_homography()
|
|
|
for kp in [12,15,16,19]:
|
|
|
self.mask_array[kp-1] = 0
|
|
|
heatmap_tensor = generate_gaussian_array_vectorized(self.num_channels, self.keypoints_final, self.size,
|
|
|
down_ratio=2, sigma=2)
|
|
|
return heatmap_tensor, self.mask_array
|
|
|
|
|
|
|
|
|
def kpmeters2yards(self, kp):
|
|
|
wp = self.keypoint_world_coords_2D[kp - 1]
|
|
|
wp_arr = np.array([wp[0] * 1.09361, wp[1] * 1.09361, 1.])
|
|
|
return wp_arr
|
|
|
|
|
|
|
|
|
|
|
|
def get_kp_from_homography(self):
|
|
|
for kp in range(1, len(self.keypoint_world_coords_2D)+1):
|
|
|
if kp not in [12, 15, 16, 19]:
|
|
|
|
|
|
wp = self.keypoint_world_coords_2D[kp-1]
|
|
|
img_pt = np.linalg.inv(self.homography) @ np.array([wp[0], wp[1], 1.])
|
|
|
img_pt /= img_pt[-1]
|
|
|
img_pt[0] *= self.w / self.image.size[0]
|
|
|
img_pt[1] *= self.h / self.image.size[1]
|
|
|
|
|
|
self.keypoints_final[kp] = {'x': img_pt[0],
|
|
|
'y': img_pt[1],
|
|
|
'in_frame': True if 0 <= img_pt[0] <= self.w and 0 <= img_pt[1] <= self.w else False,
|
|
|
'close_to_frame': True if -self.w_extra <= img_pt[0] <= self.w + self.w_extra and \
|
|
|
-self.h_extra <= img_pt[1] <= self.h + self.h_extra else False}
|
|
|
|
|
|
|
|
|
def get_lines_from_keypoints(self):
|
|
|
if len(self.keypoints_final) == 0:
|
|
|
self.get_kp_from_homography()
|
|
|
|
|
|
...
|
|
|
|
|
|
def draw_keypoints(self, scale=1):
|
|
|
|
|
|
if len(self.keypoints_final) == 0:
|
|
|
self.get_kp_from_homography()
|
|
|
|
|
|
fig, ax = plt.subplots(figsize=(scale*15, scale*7.5))
|
|
|
ax.imshow(self.image)
|
|
|
for kp in self.keypoints_final.keys():
|
|
|
if kp <= 30:
|
|
|
if self.keypoints_final[kp]['close_to_frame']:
|
|
|
x, y = self.keypoints_final[kp]['x'], self.keypoints_final[kp]['y']
|
|
|
ax.text(x, y, s=kp, zorder=11)
|
|
|
ax.scatter(x, y, c='r', s=scale*10, zorder=10)
|
|
|
|
|
|
|
|
|
elif 30 < kp <= 36:
|
|
|
if self.keypoints_final[kp]['close_to_frame']:
|
|
|
x, y = self.keypoints_final[kp]['x'], self.keypoints_final[kp]['y']
|
|
|
ax.text(x, y, s=kp, zorder=11)
|
|
|
ax.scatter(x, y, c='b', s=scale*10, zorder=10)
|
|
|
|
|
|
|
|
|
elif 36 < kp <= 44:
|
|
|
if self.keypoints_final[kp]['close_to_frame']:
|
|
|
x, y = self.keypoints_final[kp]['x'], self.keypoints_final[kp]['y']
|
|
|
ax.text(x, y, s=kp, zorder=11)
|
|
|
ax.scatter(x, y, c='pink', s=scale*10, zorder=10)
|
|
|
|
|
|
|
|
|
elif 44 < kp <= 57:
|
|
|
if self.keypoints_final[kp]['close_to_frame']:
|
|
|
x, y = self.keypoints_final[kp]['x'], self.keypoints_final[kp]['y']
|
|
|
ax.text(x, y, s=kp, zorder=11)
|
|
|
ax.scatter(x, y, c='green', s=scale*10, zorder=10)
|
|
|
|
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
