diff --git a/anipose/anipose.py b/anipose/anipose.py index 300fd6a..e7795bc 100644 --- a/anipose/anipose.py +++ b/anipose/anipose.py @@ -38,6 +38,7 @@ 'pose_2d_projected': 'pose-2d-proj', 'pose_3d': 'pose-3d', 'pose_3d_filter': 'pose-3d-filtered', + 'pose_3d_ext': 'h5', 'videos_labeled_2d': 'videos-labeled', 'videos_labeled_2d_filter': 'videos-labeled-filtered', 'calibration_videos': 'calibration', @@ -131,7 +132,7 @@ def tracking_errors(config, scorer=None): from .tracking_errors import get_tracking_errors click.echo('Comparing tracking to labeled data...') get_tracking_errors(config, scorer) - + @cli.command() @pass_config def analyze(config): @@ -236,7 +237,7 @@ def label_2d_proj(config): from .label_videos_proj import label_proj_all click.echo('Making 2D videos from 3D projections...') label_proj_all(config) - + @cli.command() @pass_config def label_2d(config): @@ -322,7 +323,7 @@ def run_data(config): from .filter_3d import filter_pose_3d_all click.echo('Filtering 3D points...') filter_pose_3d_all(config) - + click.echo('Computing angles...') compute_angles_all(config) diff --git a/anipose/triangulate.py b/anipose/triangulate.py index 97a9a10..0012925 100644 --- a/anipose/triangulate.py +++ b/anipose/triangulate.py @@ -7,6 +7,7 @@ import os.path import pandas as pd import toml +import pickle from numpy import array as arr from glob import glob from scipy import optimize @@ -81,12 +82,14 @@ def correct_coordinate_frame(config, all_points_3d, bodyparts): def load_pose2d_fnames(fname_dict, offsets_dict=None, cam_names=None): if cam_names is None: cam_names = sorted(fname_dict.keys()) + pose_names = [fname_dict[cname] for cname in cam_names] if offsets_dict is None: offsets_dict = dict([(cname, (0,0)) for cname in cam_names]) datas = [] + joints_all = [] for ix_cam, (cam_name, pose_name) in \ enumerate(zip(cam_names, pose_names)): dlabs = pd.read_hdf(pose_name) @@ -104,19 +107,21 @@ def load_pose2d_fnames(fname_dict, offsets_dict=None, cam_names=None): dlabs.loc[:, (joint, 'y')] += dy datas.append(dlabs) + joints_all += joint_names + joint_names_unique = np.unique(joints_all) n_cams = len(cam_names) - n_joints = len(joint_names) + n_joints = len(joint_names_unique) n_frames = min([d.shape[0] for d in datas]) # frame, camera, bodypart, xy points = np.full((n_cams, n_frames, n_joints, 2), np.nan, 'float') - scores = np.full((n_cams, n_frames, n_joints), np.zeros(1), 'float')#initialise as zeros, instead of NaN, makes more sense? + scores = np.full((n_cams, n_frames, n_joints), np.zeros(1), 'float')#initialise as zeros, instead of NaN, makes more sense? for cam_ix, dlabs in enumerate(datas): - for joint_ix, joint_name in enumerate(joint_names): + for joint_ix, joint_name in enumerate(joint_names_unique): try: - points[cam_ix, :, joint_ix] = np.array(dlabs.loc[:, (joint_name, ('x', 'y'))])[:n_frames] + points[cam_ix, :, joint_ix] = np.array(dlabs.loc[:, (joint_name, ('x', 'y'))])[:n_frames] scores[cam_ix, :, joint_ix] = np.array(dlabs.loc[:, (joint_name, ('likelihood'))])[:n_frames].ravel() except KeyError: pass @@ -125,7 +130,7 @@ def load_pose2d_fnames(fname_dict, offsets_dict=None, cam_names=None): 'cam_names': cam_names, 'points': points, 'scores': scores, - 'bodyparts': joint_names + 'bodyparts': joint_names_unique } @@ -290,7 +295,22 @@ def triangulate(config, dout['fnum'] = np.arange(n_frames) - dout.to_csv(output_fname, index=False) + if output_fname.endswith('.csv'): + dout.to_csv(output_fname, index=False) + else: + with open(output_fname, 'wb') as f: pickle.dump(dout, f) + + print(f'Triangulated pose is saved at: {output_fname}') + + +def get_camera_names_from_calib(calib_folder): + """ Get camera names from the calibration folder.""" + calib_fname = os.path.join(calib_folder, 'calibration.toml') + master_dict = toml.load(calib_fname) + return [ + camera_dict['name'] + for camera, camera_dict in master_dict.items() if 'cam' in camera + ] def process_session(config, session_path): @@ -299,6 +319,10 @@ def process_session(config, session_path): pipeline_pose = config['pipeline']['pose_2d'] pipeline_pose_filter = config['pipeline']['pose_2d_filter'] pipeline_3d = config['pipeline']['pose_3d'] + output_ext = config['pipeline']['pose_3d_ext'] + + if not output_ext in ['csv', 'h5', 'pkl']: + raise ValueError('The output extension should be csv, h5 or pkl!') calibration_path = find_calibration_folder(config, session_path) if calibration_path is None: @@ -313,7 +337,12 @@ def process_session(config, session_path): video_folder = os.path.join(session_path, pipeline_videos_raw) output_folder = os.path.join(session_path, pipeline_3d) - pose_files = glob(os.path.join(pose_folder, '*.h5')) + camera_names = get_camera_names_from_calib(calib_folder) + + pose_files = [] + + for cam_name in camera_names: + pose_files.append(glob(os.path.join(pose_folder, f'*{cam_name}*.h5'))[0]) cam_videos = defaultdict(list) @@ -332,14 +361,13 @@ def process_session(config, session_path): cam_names = [get_cam_name(config, f) for f in fnames] fname_dict = dict(zip(cam_names, fnames)) - output_fname = os.path.join(output_folder, name + '.csv') - + output_fname = os.path.join(output_folder, name + f'pose3d.{output_ext}') print(output_fname) - + if os.path.exists(output_fname): + print(f'Triangulation file already exists at: {output_fname}') continue - try: triangulate(config, calib_folder, video_folder, pose_folder, @@ -347,6 +375,6 @@ def process_session(config, session_path): except ValueError: import traceback, sys traceback.print_exc(file=sys.stdout) - + triangulate_all = make_process_fun(process_session)