from microprograms.temporal_segmentation.entry import entry_microprogram_one_frame from microprograms.temporal_segmentation.somersault import somersault_microprogram_one_frame from microprograms.temporal_segmentation.twist import twist_microprogram_one_frame from microprograms.temporal_segmentation.start_takeoff import takeoff_microprogram_one_frame from microprograms.errors.distance_from_springboard_micro_program import board_end from microprograms.errors.splash_micro_program import get_splash_from_one_frame from microprograms.errors.distance_from_springboard_micro_program import calculate_distance_from_springboard_for_one_frame from microprograms.errors.distance_from_springboard_micro_program import calculate_distance_from_platform_for_one_frame from microprograms.errors.angles_micro_programs import applyFeetApartError from microprograms.errors.angles_micro_programs import applyPositionTightnessError from models.detectron2.platform_detector_setup import get_platform_detector from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation from models.detectron2.diver_detector_setup import get_diver_detector from models.pose_estimator.pose_estimator_model_setup import get_pose_model from models.detectron2.splash_detector_setup import get_splash_detector from somersault_counter import som_counter, twist_counter from microprograms.errors.over_rotation import over_rotation import pickle import os import math import numpy as np import cv2 # returns None if either pose_pred or board_end_coord is None # returns True if diver is on board, returns False if diver is off board def detect_on_board(board_end_coord, board_side, pose_pred, handstand): if pose_pred is None: print("pose_pred is None") return if board_end_coord is None: print("board end coord is None") return # side = find_which_side_board_on(outputs) if board_side == 'left': # if right of board end if np.array(pose_pred)[0][2][0] > int(board_end_coord[0]): # print("board on left side, we're saying that the diver is on the right") # if standing_dive: if handstand: # print("board end to wrist dist:", math.dist(np.array(pose_pred)[0][15], board_end_coord)) # print("wrist to elbow dist:", math.dist(np.array(pose_pred)[0][14], np.array(pose_pred)[0][15])) distance = math.dist(np.array(pose_pred)[0][15], board_end_coord) < math.dist(np.array(pose_pred)[0][14], np.array(pose_pred)[0][15]) * 1.5 else: # print("board end to ankle dist:", math.dist(np.array(pose_pred)[0][5], board_end_coord)) # print("ankle to knee dist:", math.dist(np.array(pose_pred)[0][5], np.array(pose_pred)[0][4])) distance = math.dist(np.array(pose_pred)[0][5], board_end_coord) < math.dist(np.array(pose_pred)[0][5], np.array(pose_pred)[0][4]) * 1.5 # print("distance", distance) return distance # return False # if left of board end else: # print("board on left side, we're saying that the diver is on the left") return True else: # if right of board end if np.array(pose_pred)[0][2][0] > int(board_end_coord[0]): # print("board on right side, we're saying that the diver is on the right") return True # if left of board end else: # print("board on right side, we're saying that the diver is on the left") # if standing_dive: if handstand: distance = math.dist(np.array(pose_pred)[0][10], board_end_coord) < math.dist(np.array(pose_pred)[0][11], np.array(pose_pred)[0][10]) * 1.5 else: distance = math.dist(np.array(pose_pred)[0][0], board_end_coord) < math.dist(np.array(pose_pred)[0][1], np.array(pose_pred)[0][0]) * 1.5 # print("distance", distance) return distance # return False def main(): first_folder = input("what is the first folder? Ex: 01, FINAWorldChampionships2019_Women10m_final_r1, etc. ") second_folder = input("what is the second folder? (dive within the first folder)") takeoff = [] twist = [] som = [] entry = [] distance_from_board = [] position_tightness = [] feet_apart = [] splash = [] above_board = True board_side = input("what side is the board on? type either 'left' or 'right'") on_board = True handstand = input("is the dive a handstand dive? type either 'True' or 'False' ") == 'True' expected_twists = int(input("what are the expected number of twists? type 1 for half a twist, 2 for full twist, etc.")) expected_som = int(input("what are the expected number of somersaults? type 1 for half a somersault, 2 for full somersault, etc.")) # expected_twists= 0 # expected_som = 7 platform_detector = get_platform_detector() splash_detector = get_splash_detector() diver_detector = get_diver_detector() pose_model = get_pose_model() dive_folder_num = "{}_{}".format(first_folder, second_folder) directory = './FineDiving/datasets/FINADiving_MTL_256s/{}/{}/'.format(first_folder, second_folder) # dive_folder_num = '01_1' # directory = './FineDiving/datasets/FINADiving_MTL_256s/01/1/' # dive_folder_num = 'FINAWorldChampionships2019_Women10m_final_r1_0' # directory = './FineDiving/datasets/FINADiving_MTL_256s/FINAWorldChampionships2019_Women10m_final_r1/0/' file_names = os.listdir(directory) # with open('./output/joint_plots/{}/pose_preds.pkl'.format(dive_folder_num), 'rb') as pickle_file: # pose_preds = pickle.load(pickle_file) j = 0 # if len(pose_preds) > len(file_names): # print("WRONG POSE_PREDS") prev_pred = None som_prev_pred =None half_som_count=0 petal_count = 0 in_petal = False for i in range(len(file_names)): # pose_pred = None filepath = directory + file_names[i] print("filepath:", filepath) if file_names[i][-4:] != ".jpg": continue diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model) # if j < len(pose_preds): # print("filepath has pose_pred:", filepath) # pose_pred = pose_preds[j] # j += 1 calculated_half_som_count, skip = som_counter(pose_pred, som_prev_pred, half_som_count=half_som_count, handstand=handstand) if not skip: som_prev_pred = pose_pred print("calculated_half_som_count:", calculated_half_som_count) calculated_petal_count, calculated_in_petal = twist_counter(pose_pred, prev_pose_pred=prev_pred, in_petal=in_petal, petal_count=petal_count) print("calculated_petal_count", calculated_petal_count) print("calculated_in_petal", calculated_in_petal) im = cv2.imread(filepath) outputs = platform_detector(im) board_end_coord = board_end(outputs, board_side=board_side) # if board_end_coord is None: # print("NO BOARD NONE CRYING") if above_board and not on_board and board_end_coord is not None and pose_pred is not None and np.array(pose_pred)[0][2][1] > int(board_end_coord[1]): above_board=False if on_board and detect_on_board(board_end_coord, board_side, pose_pred, handstand) is not None and not detect_on_board(board_end_coord, board_side, pose_pred, handstand): on_board = False print('ON_BOARD:', on_board) print('ABOVE_BOARD:', above_board) calculated_takeoff = takeoff_microprogram_one_frame(filepath, above_board=above_board, on_board=on_board, pose_pred=pose_pred) calculated_twist = twist_microprogram_one_frame(filepath, on_board=on_board, pose_pred=pose_pred, expected_twists=expected_twists, petal_count=petal_count, expected_som=expected_som, half_som_count=half_som_count) calculated_som = somersault_microprogram_one_frame(filepath, pose_pred=pose_pred, on_board=on_board, expected_som=expected_som, half_som_count=half_som_count, expected_twists=expected_twists, petal_count=petal_count) calculated_entry = entry_microprogram_one_frame(filepath, above_board=above_board, on_board=on_board, pose_pred=pose_pred, expected_twists=expected_twists, petal_count=petal_count, expected_som=expected_som, half_som_count=half_som_count, splash_detector=splash_detector, visualize=True, dive_folder_num=dive_folder_num) if calculated_takeoff == 1: # distance from board dist = calculate_distance_from_platform_for_one_frame(filepath, visualize=False, pose_pred=pose_pred, board_end_coord=board_end_coord, platform_detector=platform_detector) # saves photo to ./output/data/distance_from_board/ distance_from_board.append(dist) # height off board # height_off_board.append(dist) # proximity to the edge of board # proximity_to_end_board.append(dist) elif calculated_som == 1: half_som_count = calculated_half_som_count # petal_count = calculated_petal_count # in_petal = calculated_in_petal if above_board: dist = calculate_distance_from_platform_for_one_frame(filepath, visualize=False, pose_pred=pose_pred, board_end_coord=board_end_coord, platform_detector=platform_detector) # saves photo to ./output/data/distance_from_board/ # distance from board distance_from_board.append(dist) # rotation speed # position tightness position_tightness.append(applyPositionTightnessError(filepath, pose_pred=pose_pred)) # feet flat # feet/legs apart feet_apart.append(applyFeetApartError(filepath, pose_pred=pose_pred)) elif calculated_twist == 1: half_som_count = calculated_half_som_count petal_count = calculated_petal_count in_petal = calculated_in_petal if above_board: dist = calculate_distance_from_platform_for_one_frame(filepath, visualize=False, pose_pred=pose_pred, board_end_coord=board_end_coord, platform_detector=platform_detector) # saves photo to ./output/data/distance_from_board/ # distance from board distance_from_board.append(dist) # twisting speed # position tightness position_tightness.append(180 - applyPositionTightnessError(filepath, pose_pred=pose_pred)) # feet flat # feet/legs apart feet_apart.append(applyFeetApartError(filepath, pose_pred=pose_pred)) elif calculated_entry == 1: # over/under twisting # over/under rotating print("OVER-ROTATION ERROR:", over_rotation(filepath, pose_pred=pose_pred, diver_detector=diver_detector, pose_model=pose_model)) # not straight during entry # splash splash.append(get_splash_from_one_frame(filepath, predictor=splash_detector, visualize=True)) else: print('no phase of dive calculated!') takeoff.append(calculated_takeoff) twist.append(calculated_twist) som.append(calculated_som) entry.append(calculated_entry) prev_pred = pose_pred print("takeoff", takeoff) print("twist", twist) print("som", som) print("entry", entry) print("distance_from_board", distance_from_board) print("position_tightness", position_tightness) print("feet_apart", feet_apart) print("splash", splash) if __name__ == "__main__": main()