NSAQA / somersault_counter.py
laurenok24's picture
Upload somersault_counter.py
8556490 verified
raw
history blame
7.06 kB
import numpy as np
import math
def distance_point_to_line_segment(px, py, x1, y1, x2, y2):
# Calculate the squared distance from point (px, py) to the line segment [(x1, y1), (x2, y2)]
def sqr_distance_point_to_segment():
line_length_sq = (x2 - x1)**2 + (y2 - y1)**2
if line_length_sq == 0:
return (px - x1)**2 + (py - y1)**2
t = max(0, min(1, ((px - x1) * (x2 - x1) + (py - y1) * (y2 - y1)) / line_length_sq))
return ((px - (x1 + t * (x2 - x1)))**2 + (py - (y1 + t * (y2 - y1)))**2)
# Calculate the closest point on the line segment to the given point (px, py)
def closest_point_on_line_segment():
line_length_sq = (x2 - x1)**2 + (y2 - y1)**2
if line_length_sq == 0:
return x1, y1
t = max(0, min(1, ((px - x1) * (x2 - x1) + (py - y1) * (y2 - y1)) / line_length_sq))
closest_x = x1 + t * (x2 - x1)
closest_y = y1 + t * (y2 - y1)
return closest_x, closest_y
closest_point = closest_point_on_line_segment()
distance = math.sqrt(sqr_distance_point_to_segment())
return closest_point, distance
def min_distance_from_line_to_circle(line_start, line_end, circle_center, circle_radius):
closest_point, distance = distance_point_to_line_segment(circle_center[0], circle_center[1],
line_start[0], line_start[1],
line_end[0], line_end[1])
min_distance = max(0, distance - circle_radius)
return min_distance
# def twist_counter(pose_pred, prev_pose_pred=None, in_petal=False, petal_count=0):
# if pose_pred is None:
# return petal_count, in_petal
# min_dist = 0
# # Users/lokamoto/Comprehensive_AQA/output/joint_plots/FINAWorldChampionships2019_Women10m_final_r1_0
# pose_pred = pose_pred[0]
# vector1 = [pose_pred[2][0] - pose_pred[3][0], 0-(pose_pred[2][1] - pose_pred[3][1])]
# if prev_pose_pred is not None:
# prev_pose_pred = prev_pose_pred[0]
# prev_pose_pred = [prev_pose_pred[2][0] - prev_pose_pred[3][0], 0-(prev_pose_pred[2][1] - prev_pose_pred[3][1])]
# # m = (vector1[1] - prev_pose_pred[1])/(vector1[0] - prev_pose_pred[0])
# # b = prev_pose_pred[1] - m * prev_pose_pred[0]
# # min_dist = np.abs(b)/np.sqrt(m**2+1)
# min_dist = min_distance_from_line_to_circle(prev_pose_pred, vector1, (0, 0), 5)
# if min_dist is not None and in_petal and np.linalg.norm(vector1) > 5 and min_dist == 0: #and np.linalg.norm(vector1) > 8
# petal_count += 1
# # print('leaving petal')
# # print('going in new petal')
# elif not in_petal and np.linalg.norm(vector1) > 5: #and min_dist > 3: #and np.linalg.norm(vector1) > 8
# in_petal = True
# # print('going in petal')
# elif in_petal and np.linalg.norm(vector1) < 5:
# in_petal = False
# petal_count += 1
# # print('leaving petal')
# # print(vector)
# return petal_count, in_petal
def twist_counter(pose_pred, prev_pose_pred=None, in_petal=False, petal_count=0):
# if key[0].startswith('FINAWorldChampionships2019'):
# valid = 19
# outer = 3
# inner = 3
# elif key[0].startswith('FINADivingWorldCup2021'):
# valid = 19
# outer = 5.5
# inner = 5.5
valid = 17
outer = 10
inner = 9
if pose_pred is None:
return petal_count, in_petal
min_dist = 0
pose_pred = pose_pred[0]
vector1 = [pose_pred[2][0] - pose_pred[3][0], 0-(pose_pred[2][1] - pose_pred[3][1])]
if prev_pose_pred is not None:
prev_pose_pred = prev_pose_pred[0]
prev_pose_pred = [prev_pose_pred[2][0] - prev_pose_pred[3][0], 0-(prev_pose_pred[2][1] - prev_pose_pred[3][1])]
min_dist = min_distance_from_line_to_circle(prev_pose_pred, vector1, (0, 0), 0.5)
if np.linalg.norm(vector1) > valid:
return petal_count, in_petal
if min_dist is not None and in_petal and np.linalg.norm(vector1) > outer and min_dist == 0: #and np.linalg.norm(vector1) > 8
petal_count += 1
elif not in_petal and np.linalg.norm(vector1) > outer: #and min_dist > 3: #and np.linalg.norm(vector1) > 8
in_petal = True
elif in_petal and np.linalg.norm(vector1) < inner:
in_petal = False
petal_count += 1
return petal_count, in_petal
# def som_counter(pose_pred=None, prev_pose_pred=None, half_som_count=0, handstand=False):
# if pose_pred is None:
# return half_som_count
# pose_pred = pose_pred[0]
# vector1 = [pose_pred[7][0] - pose_pred[6][0], 0-(pose_pred[7][1] - pose_pred[6][1])] # flip y axis
# if (not handstand and half_som_count % 2 == 0) or (handstand and half_som_count %2 == 1):
# vector2 = [0, -1]
# else:
# vector2 = [0, 1]
# unit_vector_1 = vector1 / np.linalg.norm(vector1)
# unit_vector_2 = vector2 / np.linalg.norm(vector2)
# dot_product = np.dot(unit_vector_1, unit_vector_2)
# current_angle = math.degrees(np.arccos(dot_product))
# if prev_pose_pred is not None:
# prev_pose_pred = prev_pose_pred[0]
# prev_vector = [prev_pose_pred[7][0] - prev_pose_pred[6][0], 0-(prev_pose_pred[7][1] - prev_pose_pred[6][1])] # flip y axis
# prev_unit_vector = prev_vector / np.linalg.norm(prev_vector)
# prev_angle_diff = math.degrees(np.arccos(np.dot(unit_vector_1, prev_unit_vector)))
# # if prev_angle_diff > 120:
# # print('pose pred is probably off')
# # return half_som_count
# # print("unit_vector_1:", unit_vector_1)
# # print("looking for vector:", vector2)
# if current_angle < 80:
# half_som_count += 1
# return half_som_count
def som_counter(pose_pred=None, prev_pose_pred=None, half_som_count=0, handstand=False):
if pose_pred is None:
return half_som_count, True
pose_pred = pose_pred[0]
vector1 = [pose_pred[7][0] - pose_pred[6][0], 0-(pose_pred[7][1] - pose_pred[6][1])] # flip y axis
if (not handstand and half_som_count % 2 == 0) or (handstand and half_som_count %2 == 1):
vector2 = [0, -1]
else:
vector2 = [0, 1]
unit_vector_1 = vector1 / np.linalg.norm(vector1)
unit_vector_2 = vector2 / np.linalg.norm(vector2)
dot_product = np.dot(unit_vector_1, unit_vector_2)
current_angle = math.degrees(np.arccos(dot_product))
if prev_pose_pred is not None:
prev_pose_pred = prev_pose_pred[0]
prev_vector = [prev_pose_pred[7][0] - prev_pose_pred[6][0], 0-(prev_pose_pred[7][1] - prev_pose_pred[6][1])] # flip y axis
prev_unit_vector = prev_vector / np.linalg.norm(prev_vector)
prev_angle_diff = math.degrees(np.arccos(np.dot(unit_vector_1, prev_unit_vector)))
if prev_angle_diff > 115:
return half_som_count, True
# print("unit_vector_1:", unit_vector_1)
# print("looking for vector:", vector2)
if current_angle <= 80:
half_som_count += 1
return half_som_count, False