diff --git a/Pose2Sim/Pose2Sim.py b/Pose2Sim/Pose2Sim.py index ab0479f..4abaf8b 100644 --- a/Pose2Sim/Pose2Sim.py +++ b/Pose2Sim/Pose2Sim.py @@ -345,6 +345,8 @@ def triangulation(config=None): ''' from Pose2Sim.triangulation import triangulate_all + from Pose2Sim.augmenter import augmentTRC + from Pose2Sim.filtering import filter_all # Determine the level at which the function is called (session:3, participant:2, trial:1) level, config_dicts = read_config_files(config) @@ -371,8 +373,10 @@ def triangulation(config=None): logging.info(f"Triangulation of 2D points for {seq_name}, for {frames}.") logging.info("---------------------------------------------------------------------") logging.info(f"\nProject directory: {project_dir}") - - triangulate_all(config_dict) + if config_dict.get('pose').get('pose_model') == 'BODY_25_AUGMENTED': + triangulation_filt_AUGMENTED(config_dict) + else: + triangulate_all(config_dict) end = time.time() logging.info(f'Triangulation took {end-start:.2f} s.') @@ -549,6 +553,7 @@ def triangulation_filt_AUGMENTED(config=None): end = time.time() logging.info(f'Augmentation took {end - start:.2f} s.') + def filtering(config=None): diff --git a/Pose2Sim/filtering.py b/Pose2Sim/filtering.py index 1917e20..7dd6a2e 100644 --- a/Pose2Sim/filtering.py +++ b/Pose2Sim/filtering.py @@ -464,9 +464,17 @@ def filter_all(config): # Trc paths trc_f_in = f'{seq_name}_{f_range[0]}-{f_range[1]}.trc' + trc_f_in_LSTM = f'{seq_name}_{f_range[0]}-{f_range[1]}_LSTM.trc' trc_f_out = f'{seq_name}_filt_{filter_type}_{f_range[0]}-{f_range[1]}.trc' - trc_path_in = os.path.join(pose3d_dir, trc_f_in) - trc_path_out = os.path.join(pose3d_dir, trc_f_out) + trc_f_out_LSTM = f'{seq_name}_filt_{filter_type}_{f_range[0]}-{f_range[1]}_LSTM.trc' + if os.path.exists(os.path.join(pose3d_dir, trc_f_in_LSTM)): + trc_path_in = os.path.join(pose3d_dir, trc_f_in_LSTM) + trc_path_out = os.path.join(pose3d_dir, trc_f_out_LSTM) + elif os.path.exists(os.path.join(pose3d_dir, trc_f_in)): + trc_path_in = os.path.join(pose3d_dir, trc_f_in) + trc_path_out = os.path.join(pose3d_dir, trc_f_out) + else: + raise FileNotFoundError("Not found .trc file.") # Read trc header with open(trc_path_in, 'r') as trc_file: diff --git a/Pose2Sim/skeletons.py b/Pose2Sim/skeletons.py index 6ec0244..5792c64 100644 --- a/Pose2Sim/skeletons.py +++ b/Pose2Sim/skeletons.py @@ -77,6 +77,46 @@ BODY_25B = Node("CHip", id=None, children=[ ]), ]) +'''BODY_25B_AUGMENTED (full-body without hands, experimental, from OpenPose) +https://github.com/CMU-Perceptual-Computing-Lab/openpose_train/blob/master/experimental_models/README.md''' +BODY_25B_AUGMENTED = Node("CHip", id=None, children=[ + Node("RHip", id=12, children=[ + Node("RKnee", id=14, children=[ + Node("RAnkle", id=16, children=[ + Node("RBigToe", id=22, children=[ + Node("RSmallToe", id=23), + ]), + Node("RHeel", id=24), + ]), + ]), + ]), + Node("LHip", id=11, children=[ + Node("LKnee", id=13, children=[ + Node("LAnkle", id=15, children=[ + Node("LBigToe", id=19, children=[ + Node("LSmallToe", id=20), + ]), + Node("LHeel", id=21), + ]), + ]), + ]), + Node("Neck", id=17, children=[ + Node("Head", id=18, children=[ + Node("Nose", id=0), + ]), + Node("RShoulder", id=6, children=[ + Node("RElbow", id=8, children=[ + Node("RWrist", id=10), + ]), + ]), + Node("LShoulder", id=5, children=[ + Node("LElbow", id=7, children=[ + Node("LWrist", id=9), + ]), + ]), + ]), +]) + '''BODY_25 (full-body without hands, standard, from OpenPose) https://github.com/CMU-Perceptual-Computing-Lab/openpose/tree/master/models''' @@ -116,6 +156,44 @@ BODY_25 = Node("CHip", id=8, children=[ ]), ]) +'''BODY_25_AUGMENTED (full-body without hands, standard, from OpenPose) +https://github.com/CMU-Perceptual-Computing-Lab/openpose/tree/master/models''' +BODY_25_AUGMENTED = Node("CHip", id=8, children=[ + Node("RHip", id=9, children=[ + Node("RKnee", id=10, children=[ + Node("RAnkle", id=11, children=[ + Node("RBigToe", id=22, children=[ + Node("RSmallToe", id=23), + ]), + Node("RHeel", id=24), + ]), + ]), + ]), + Node("LHip", id=12, children=[ + Node("LKnee", id=13, children=[ + Node("LAnkle", id=14, children=[ + Node("LBigToe", id=19, children=[ + Node("LSmallToe", id=20), + ]), + Node("LHeel", id=21), + ]), + ]), + ]), + Node("Neck", id=17, children=[ + Node("Nose", id=0), + Node("RShoulder", id=2, children=[ + Node("RElbow", id=3, children=[ + Node("RWrist", id=4), + ]), + ]), + Node("LShoulder", id=5, children=[ + Node("LElbow", id=6, children=[ + Node("LWrist", id=7), + ]), + ]), + ]), +]) + '''BODY_135 (full-body with hands and face, experimental, from OpenPose) https://github.com/CMU-Perceptual-Computing-Lab/openpose_train/blob/master/experimental_models/README.md)'''