Upload somersault_counter.py

somersault_counter.py

@@ -0,0 +1,153 @@
+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
+