diff --git a/.gitignore b/.gitignore index c812d87..c8b587a 100644 --- a/.gitignore +++ b/.gitignore @@ -24,4 +24,6 @@ logs !.vscode/extensions.json *.code-workspace -# End of https://www.toptal.com/developers/gitignore/api/vscode \ No newline at end of file +# End of https://www.toptal.com/developers/gitignore/api/vscode + +src/*.sh \ No newline at end of file diff --git a/conda_envs/acinoset.yml b/conda_envs/acinoset.yml index b96a8ee..319478b 100644 --- a/conda_envs/acinoset.yml +++ b/conda_envs/acinoset.yml @@ -22,3 +22,4 @@ dependencies: - pyqt5 - tables - nptyping==0.3.1 + - tqdm diff --git a/docker/Dockerfile.acinoset b/docker/Dockerfile.acinoset.gui similarity index 100% rename from docker/Dockerfile.acinoset rename to docker/Dockerfile.acinoset.gui diff --git a/run_acinoset.sh b/run_acinoset.sh old mode 100644 new mode 100755 index 17238a5..7271cc6 --- a/run_acinoset.sh +++ b/run_acinoset.sh @@ -3,17 +3,10 @@ CURRENT_PATH=$(pwd) IMAGE_NAME="denden047/acinoset" -docker build -q -f "$CURRENT_PATH"/docker/Dockerfile.acinoset -t ${IMAGE_NAME} . && \ +docker build -q -f "$CURRENT_PATH"/docker/Dockerfile.acinoset.cui -t ${IMAGE_NAME} . && \ docker run -it --rm \ - --gpus device=0 \ -v "$CURRENT_PATH":/workdir \ - -v /disk2/naoya/AcinoSet:/data \ + -v /data/naoya/AcinoSet:/data \ -w /workdir \ - -p 6080:80 \ ${IMAGE_NAME} \ /bin/bash - # /bin/bash -c " \ - # conda init bash && \ - # conda activate acinoset && \ - # jupyter lab --allow-root --NotebookApp.token='' --ip=* --no-browser \ - # " \ No newline at end of file diff --git a/src/all_flick.sh b/src/all_flick.sh index a04d2f8..a49c5f1 100644 --- a/src/all_flick.sh +++ b/src/all_flick.sh @@ -9,7 +9,7 @@ python all_optimizations.py --data_dir /data/2019_03_09/jules/flick1 python all_optimizations.py --data_dir /data/2019_03_09/jules/flick2 python all_optimizations.py --data_dir /data/2019_03_03/menya/flick python all_optimizations.py --data_dir /data/2019_03_03/phantom/flick -python all_optimizations.py --data_dir /data/2019_02_27/Kiara/flick +python all_optimizations.py --data_dir /data/2019_02_27/kiara/flick python all_optimizations.py --data_dir /data/2019_02_27/romeo/flick python all_optimizations.py --data_dir /data/2017_08_29/bottom/phantom/flick2 python all_optimizations.py --data_dir /data/2017_08_29/bottom/zorro/flick2 diff --git a/src/all_optimizations.py b/src/all_optimizations.py index 3706d4a..a7e3314 100644 --- a/src/all_optimizations.py +++ b/src/all_optimizations.py @@ -19,7 +19,7 @@ plt.style.use(os.path.join('..', 'configs', 'mplstyle.yaml')) -def fte(DATA_DIR, start_frame, end_frame, dlc_thresh): +def fte(DATA_DIR, DLC_DIR, start_frame, end_frame, dlc_thresh, plot: bool = False): # PLOT OF REDESCENDING, ABSOLUTE AND QUADRATIC COST FUNCTIONS # we use a redescending cost to stop outliers affecting the optimisation negatively redesc_a = 3 @@ -45,8 +45,6 @@ def fte(DATA_DIR, start_frame, end_frame, dlc_thresh): assert os.path.exists(DATA_DIR) OUT_DIR = os.path.join(DATA_DIR, 'fte') - DLC_DIR = os.path.join(DATA_DIR, 'dlc') - assert os.path.exists(DLC_DIR) os.makedirs(OUT_DIR, exist_ok=True) app.start_logging(os.path.join(OUT_DIR, 'fte.log')) @@ -178,9 +176,9 @@ def rot_z(z): # p_lure.T ]) - func_map = {"sin":sin, "cos":cos, "ImmutableDenseMatrix":np.array} + func_map = {"sin":sin, "cos":cos, "ImmutableDenseMatrix":np.array} sym_list = [x, y, z, - *phi, *theta, *psi, + *phi, *theta, *psi, # x_l, y_l, z_l ] pose_to_3d = sp.lambdify(sym_list, positions, modules=[func_map]) @@ -260,7 +258,7 @@ def get_likelihood_from_df(n, c, l): points_2d_df = utils.load_dlc_points_as_df(df_paths, verbose=False) points_3d_df = utils.get_pairwise_3d_points_from_df( - points_2d_df[points_2d_df['likelihood']>dlc_thresh], + points_2d_df[points_2d_df['likelihood'] > dlc_thresh], K_arr, D_arr, R_arr, t_arr, triangulate_points_fisheye ) @@ -447,7 +445,7 @@ def tail_mid_theta_5(m,n): return abs(m.x[n,23]) <= np.pi/1.5 m.tail_mid_theta_5 = Constraint(m.N, rule=tail_mid_theta_5) def tail_mid_psi_5(m,n): - return abs(m.x[n,37]) <= np.pi/1.5 + return abs(m.x[n,37]) <= np.pi/1.5 m.tail_mid_psi_5 = Constraint(m.N, rule=tail_mid_psi_5) #Front left leg @@ -554,7 +552,7 @@ def convert_m(m, pose_indices): [unused_pose_indices] = np.where(Q == 0) positions, states = convert_m(m, unused_pose_indices) - + out_fpath = os.path.join(OUT_DIR, f"fte.pickle") app.save_optimised_cheetah(positions, out_fpath, extra_data=dict(**states, start_frame=start_frame)) app.save_3d_cheetah_as_2d(positions, OUT_DIR, scene_fpath, markers, project_points_fisheye, start_frame) @@ -564,17 +562,14 @@ def convert_m(m, pose_indices): fig_fpath= os.path.join(OUT_DIR, 'fte.svg') app.plot_cheetah_states(states['x'], out_fpath=fig_fpath) - - -def ekf(DATA_DIR, start_frame, end_frame, dlc_thresh): + + +def ekf(DATA_DIR, DLC_DIR, start_frame, end_frame, dlc_thresh): # ========= INIT VARS ======== t0 = time() - assert os.path.exists(DATA_DIR) OUT_DIR = os.path.join(DATA_DIR, 'ekf') - DLC_DIR = os.path.join(DATA_DIR, 'dlc') - assert os.path.exists(DLC_DIR) os.makedirs(OUT_DIR, exist_ok=True) app.start_logging(os.path.join(OUT_DIR, 'ekf.log')) @@ -685,7 +680,7 @@ def numerical_jacobian(func, x: np.ndarray, *args): # try: # lure_pts = points_3d_df[points_3d_df["marker"]=="lure"][["frame", "x", "y", "z"]].values - # lure_x_slope, lure_x_intercept, *_ = linregress(lure_pts[:,0], lure_pts[:,1]) + # lure_x_slope, lure_x_intercept, *_ = linregress(lure_pts[:,0], lure_pts[:,1]) # lure_y_slope, lure_y_intercept, *_ = linregress(lure_pts[:,0], lure_pts[:,2]) # lure_x_est = start_frame*lure_x_slope + lure_x_intercept # initial lure x @@ -865,47 +860,66 @@ def numerical_jacobian(func, x: np.ndarray, *args): app.plot_cheetah_states(states['x'], states['smoothed_x'], fig_fpath) -def sba(DATA_DIR: str, DLC_DIR: str, scene_fpath: str, num_frame: int, dlc_thresh: float, plot: bool = False): - # init +def sba(DATA_DIR, DLC_DIR, start_frame, end_frame, dlc_thresh): + t0 = time() + OUT_DIR = os.path.join(DATA_DIR, 'sba') os.makedirs(OUT_DIR, exist_ok=True) + app.start_logging(os.path.join(OUT_DIR, 'sba.log')) + # load video info + res, fps, tot_frames, _ = app.get_vid_info(DATA_DIR) # path to original videos + assert end_frame <= tot_frames, f'end_frame must be less than or equal to {tot_frames}' + + start_frame -= 1 # 0 based indexing + assert start_frame >= 0 + N = end_frame-start_frame + + *_, n_cams, scene_fpath = utils.find_scene_file(DATA_DIR, verbose=False) + + dlc_points_fpaths = sorted(glob(os.path.join(DLC_DIR, '*.h5'))) + assert n_cams == len(dlc_points_fpaths) + + # Load Measurement Data (pixels, likelihood) + points_2d_df = utils.load_dlc_points_as_df(dlc_points_fpaths, verbose=False) + points_2d_df = points_2d_df[points_2d_df["frame"].between(start_frame, end_frame-1)] + points_2d_df = points_2d_df[points_2d_df['likelihood']>dlc_thresh] # ignore points with low likelihood + + t1 = time() + print("Initialization took {0:.2f} seconds\n".format(t1 - t0)) + points_3d_df, residuals = app.sba_points_fisheye(scene_fpath, points_2d_df) app.stop_logging() - plt.plot(residuals['before'], alpha=0.5, label="Cost before") - plt.plot(residuals['after'], alpha=0.5, label="Cost after") + plt.plot(residuals['before'], label="Cost before") + plt.plot(residuals['after'], label="Cost after") plt.legend() - fig_fpath = os.path.join(OUT_DIR, 'sba.pdf') + fig_fpath = os.path.join(OUT_DIR, 'sba.svg') plt.savefig(fig_fpath, transparent=True) print(f'Saved {fig_fpath}\n') - if plot: - plt.show(block=True) + plt.show(block=False) # ========= SAVE SBA RESULTS ======== + markers = misc.get_markers() - positions = np.full((num_frame, len(markers), 3), np.nan) + + positions = np.full((N, len(markers), 3), np.nan) for i, marker in enumerate(markers): marker_pts = points_3d_df[points_3d_df["marker"]==marker][["frame", "x", "y", "z"]].values for frame, *pt_3d in marker_pts: - positions[int(frame) - start_frame, i] = pt_3d + positions[int(frame)-start_frame, i] = pt_3d app.save_sba(positions, OUT_DIR, scene_fpath, start_frame, dlc_thresh) - - -def tri(DATA_DIR, start_frame, end_frame, dlc_thresh): - assert os.path.exists(DATA_DIR) + + +def tri(DATA_DIR, DLC_DIR, start_frame: int, end_frame: int, dlc_thresh): OUT_DIR = os.path.join(DATA_DIR, 'tri') - DLC_DIR = os.path.join(DATA_DIR, 'dlc') - assert os.path.exists(DLC_DIR) os.makedirs(OUT_DIR, exist_ok=True) - -def tri(DATA_DIR: str, DLC_DIR: str, start_frame: int, end_frame: int, dlc_thresh: float): - out_dir = os.path.join(DATA_DIR, 'tri') - os.makedirs(out_dir, exist_ok=True) + start_frame -= 1 # 0 based indexing + assert start_frame >= 0 N = end_frame - start_frame k_arr, d_arr, r_arr, t_arr, cam_res, n_cams, scene_fpath = utils.find_scene_file(DATA_DIR, verbose=False) @@ -916,9 +930,9 @@ def tri(DATA_DIR: str, DLC_DIR: str, start_frame: int, end_frame: int, dlc_thres # Load Measurement Data (pixels, likelihood) points_2d_df = utils.load_dlc_points_as_df(dlc_points_fpaths, verbose=False) points_2d_df = points_2d_df[points_2d_df["frame"].between(start_frame, end_frame-1)] - points_2d_df = points_2d_df[points_2d_df['likelihood'] > dlc_thresh] # ignore points with low likelihood + points_2d_df = points_2d_df[points_2d_df['likelihood']>dlc_thresh] # ignore points with low likelihood - assert len(k_arr) == points_2d_df['camera'].nunique() + assert len(k_arr) == points_2d_df['camera'].nunique(), '{}'.format(points_2d_df['camera']) points_3d_df = utils.get_pairwise_3d_points_from_df( points_2d_df, @@ -929,22 +943,24 @@ def tri(DATA_DIR: str, DLC_DIR: str, start_frame: int, end_frame: int, dlc_thres points_3d_df['point_index'] = points_3d_df.index # ========= SAVE TRIANGULATION RESULTS ======== + markers = misc.get_markers() + positions = np.full((N, len(markers), 3), np.nan) for i, marker in enumerate(markers): marker_pts = points_3d_df[points_3d_df["marker"]==marker][["frame", "x", "y", "z"]].values for frame, *pt_3d in marker_pts: - positions[int(frame) - start_frame, i] = pt_3d + positions[int(frame)-start_frame, i] = pt_3d app.save_tri(positions, OUT_DIR, scene_fpath, start_frame, dlc_thresh) - - + + def dlc(DATA_DIR, dlc_thresh): video_fpaths = sorted(glob(os.path.join(DATA_DIR, 'cam[1-9].mp4'))) # original vids should be in the parent dir out_dir = os.path.join(DATA_DIR, 'dlc') app.create_labeled_videos(video_fpaths, out_dir=out_dir, draw_skeleton=True, pcutoff=dlc_thresh) - - + + # ========= MAIN ======== if __name__ == "__main__": @@ -956,29 +972,39 @@ def dlc(DATA_DIR, dlc_thresh): parser.add_argument('--dlc_thresh', type=float, default=0.8, help='The likelihood of the dlc points below which will be excluded from the optimization') parser.add_argument('--plot', action='store_true', help='Showing plots') args = parser.parse_args() - - ROOT_DATA_DIR = os.path.join("..", "data") - DATA_DIR = os.path.join(ROOT_DATA_DIR, os.path.normpath(args.data_dir)) - + + # ROOT_DATA_DIR = os.path.join("..", "data") + DATA_DIR = os.path.normpath(args.data_dir) + assert os.path.exists(DATA_DIR) + DLC_DIR = os.path.join(DATA_DIR, 'dlc') + assert os.path.exists(DLC_DIR) + + # load DLC info + res, fps, tot_frames, _ = app.get_vid_info(DATA_DIR) # path to original videos + assert args.end_frame <= tot_frames, f'end_frame must be less than or equal to {tot_frames}' + end_frame = tot_frames if args.end_frame == -1 else args.end_frame + print('========== DLC ==========\n') dlc(DATA_DIR, args.dlc_thresh) print('========== Triangulation ==========\n') - tri(DATA_DIR, args.start_frame, args.end_frame, args.dlc_thresh) + tri(DATA_DIR, DLC_DIR, args.start_frame, end_frame, args.dlc_thresh) plt.close('all') print('========== SBA ==========\n') - sba(DATA_DIR, args.start_frame, args.end_frame, args.dlc_thresh) + sba(DATA_DIR, DLC_DIR, args.start_frame, end_frame, args.dlc_thresh) plt.close('all') print('========== EKF ==========\n') - ekf(DATA_DIR, args.start_frame, args.end_frame, args.dlc_thresh) + ekf(DATA_DIR, DLC_DIR, args.start_frame, end_frame, args.dlc_thresh) plt.close('all') print('========== FTE ==========\n') - fte(DATA_DIR, args.start_frame, args.end_frame, args.dlc_thresh) + fte(DATA_DIR, DLC_DIR, args.start_frame, end_frame, args.dlc_thresh) plt.close('all') - - print('Plotting results...') - data_fpaths = [#os.path.join(DATA_DIR, 'tri', 'tri.pickle'), # plot is too busy when tri is included - os.path.join(DATA_DIR, 'sba', 'sba.pickle'), - os.path.join(DATA_DIR, 'ekf', 'ekf.pickle'), - os.path.join(DATA_DIR, 'fte', 'fte.pickle')] - app.plot_multiple_cheetah_reconstructions(data_fpaths, reprojections=False, dark_mode=True) - \ No newline at end of file + + if args.plot: + print('Plotting results...') + data_fpaths = [ + os.path.join(DATA_DIR, 'tri', 'tri.pickle'), # plot is too busy when tri is included + os.path.join(DATA_DIR, 'sba', 'sba.pickle'), + os.path.join(DATA_DIR, 'ekf', 'ekf.pickle'), + os.path.join(DATA_DIR, 'fte', 'fte.pickle') + ] + app.plot_multiple_cheetah_reconstructions(data_fpaths, reprojections=False, dark_mode=True) \ No newline at end of file diff --git a/src/lib/app.py b/src/lib/app.py index 1479224..1a410a2 100644 --- a/src/lib/app.py +++ b/src/lib/app.py @@ -34,7 +34,7 @@ def extract_corners_from_images(img_dir, out_fpath, board_shape, board_edge_len, os.remove(f) save_points(out_fpath, saved_points, saved_fnames, board_shape, board_edge_len, cam_res) - + # ========== CALIBRATION ========== def calibrate_standard_intrinsics(points_fpath, out_fpath): @@ -64,7 +64,7 @@ def calibrate_fisheye_extrinsics_pairwise(camera_fpaths, points_fpaths, out_fpat # ========== SBA ========== - + def sba_board_points_standard(scene_fpath, points_fpaths, out_fpath, manual_points_fpath=None, manual_points_only=False, camera_indices=None): triangulate_func = triangulate_points project_func = project_points @@ -81,14 +81,14 @@ def sba_points_standard(scene_fpath, points_2d_df): triangulate_func = triangulate_points project_func = project_points return _sba_points(scene_fpath, points_2d_df, triangulate_func, project_func) - + def sba_points_fisheye(scene_fpath, points_2d_df): triangulate_func = triangulate_points_fisheye project_func = project_points_fisheye return _sba_points(scene_fpath, points_2d_df, triangulate_func, project_func) - + # ========== PLOTTING ========== def plot_corners(points_fpath): @@ -110,8 +110,8 @@ def plot_points_fisheye_undistort(points_fpath, camera_fpath): undistort_pts = create_undistort_fisheye_point_function(k, d) undistorted_points = undistort_pts(points).reshape(points.shape) plot_calib_board(undistorted_points, board_shape, cam_res) - - + + def plot_scene(data_dir, scene_fname=None, manual_points_only=False, **kwargs): *_, scene_fpath = find_scene_file(data_dir, scene_fname, verbose=False) points_dir = os.path.join(os.path.dirname(scene_fpath), "points") @@ -129,29 +129,29 @@ def plot_scene(data_dir, scene_fname=None, manual_points_only=False, **kwargs): frames.append(img_names) plot_extrinsics(scene_fpath, pts_2d, frames, triangulate_points_fisheye, manual_points_only, **kwargs) - - + + def plot_cheetah_states(states, smoothed_states=None, out_fpath=None, mplstyle_fpath=None): fig, axs = plot_optimized_states(states, smoothed_states, mplstyle_fpath) if out_fpath is not None: fig.savefig(out_fpath, transparent=True) print(f'Saved {out_fpath}\n') - + def _plot_cheetah_reconstruction(positions, data_dir, scene_fname=None, labels=None, **kwargs): positions = np.array(positions) *_, scene_fpath = find_scene_file(data_dir, scene_fname, verbose=False) ca = Cheetah(positions, scene_fpath, labels, project_points_fisheye, **kwargs) ca.animation() - - + + def plot_cheetah_reconstruction(data_fpath, scene_fname=None, **kwargs): label = os.path.basename(os.path.splitext(data_fpath)[0]).upper() with open(data_fpath, 'rb') as f: data = pickle.load(f) positions = data["smoothed_positions"] if 'EKF' in label else data["positions"] _plot_cheetah_reconstruction([positions], os.path.dirname(data_fpath), scene_fname, labels=[label], **kwargs) - + def plot_multiple_cheetah_reconstructions(data_fpaths, scene_fname=None, **kwargs): positions = [] @@ -173,17 +173,17 @@ def save_tri(positions, out_dir, scene_fpath, start_frame, dlc_thresh, save_vide out_fpath = os.path.join(out_dir, f"tri.pickle") save_optimised_cheetah(positions, out_fpath, extra_data=dict(start_frame=start_frame)) save_3d_cheetah_as_2d(positions, out_dir, scene_fpath, get_markers(), project_points_fisheye, start_frame) - + if save_videos: video_fpaths = sorted(glob(os.path.join(os.path.dirname(out_dir), 'cam[1-9].mp4'))) # original vids should be in the parent dir create_labeled_videos(video_fpaths, out_dir=out_dir, draw_skeleton=True, pcutoff=dlc_thresh) - + def save_sba(positions, out_dir, scene_fpath, start_frame, dlc_thresh, save_videos=True): out_fpath = os.path.join(out_dir, f"sba.pickle") save_optimised_cheetah(positions, out_fpath, extra_data=dict(start_frame=start_frame)) save_3d_cheetah_as_2d(positions, out_dir, scene_fpath, get_markers(), project_points_fisheye, start_frame) - + if save_videos: video_fpaths = sorted(glob(os.path.join(os.path.dirname(out_dir), 'cam[1-9].mp4'))) # original vids should be in the parent dir create_labeled_videos(video_fpaths, out_dir=out_dir, draw_skeleton=True, pcutoff=dlc_thresh) @@ -196,7 +196,7 @@ def save_ekf(states, out_dir, scene_fpath, start_frame, dlc_thresh, save_videos= out_fpath = os.path.join(out_dir, f"ekf.pickle") save_optimised_cheetah(positions, out_fpath, extra_data=dict(smoothed_positions=smoothed_positions, **states, start_frame=start_frame)) save_3d_cheetah_as_2d(smoothed_positions, out_dir, scene_fpath, get_markers(), project_points_fisheye, start_frame) - + if save_videos: video_fpaths = sorted(glob(os.path.join(os.path.dirname(out_dir), 'cam[1-9].mp4'))) # original vids should be in the parent dir create_labeled_videos(video_fpaths, out_dir=out_dir, draw_skeleton=True, pcutoff=dlc_thresh) @@ -207,19 +207,19 @@ def save_fte(states, out_dir, scene_fpath, start_frame, dlc_thresh, save_videos= out_fpath = os.path.join(out_dir, f"fte.pickle") save_optimised_cheetah(positions, out_fpath, extra_data=dict(**states, start_frame=start_frame)) save_3d_cheetah_as_2d(positions, out_dir, scene_fpath, get_markers(), project_points_fisheye, start_frame) - + if save_videos: video_fpaths = sorted(glob(os.path.join(os.path.dirname(out_dir), 'cam[1-9].mp4'))) # original vids should be in the parent dir create_labeled_videos(video_fpaths, out_dir=out_dir, draw_skeleton=True, pcutoff=dlc_thresh) - + # ========== STDOUT LOGGING ========== def start_logging(out_fpath): """Start logger, appending print output to given output file""" sys.stdout = Logger(out_fpath) - + def stop_logging(): """Stop logging and return print functionality to normal""" sys.stdout.logfile.close() @@ -230,11 +230,11 @@ def stop_logging(): def get_vid_info(path_dir, vid_extension='mp4'): """Finds a video specified in/by the path variable and returns its info - + :param path: Either a directory containing a video or the path to a specific video file """ from errno import ENOENT - + orig_path = path_dir if not os.path.isfile(path_dir): files = sorted(glob(os.path.join(path_dir, f"*.{vid_extension}"))) # assume path is a dir that holds video file(s) @@ -242,7 +242,7 @@ def get_vid_info(path_dir, vid_extension='mp4'): path_dir = files[0] else: raise FileNotFoundError(ENOENT, os.strerror(ENOENT), orig_path) # assume videos didn't open due to incorrect path - + vid = VideoProcessorCV(in_name=path_dir) vid.close() return (vid.width(), vid.height()), vid.fps(), vid.frame_count(), vid.codec() @@ -253,7 +253,7 @@ def create_labeled_videos(video_fpaths, videotype="mp4", codec="mp4v", outputfra from multiprocessing import Pool print('Saving labeled videos...') - + bodyparts = get_markers() bodyparts2connect = get_skeleton() if draw_skeleton else None @@ -268,5 +268,5 @@ def create_labeled_videos(video_fpaths, videotype="mp4", codec="mp4v", outputfra with Pool(min(os.cpu_count(), len(video_fpaths))) as pool: pool.map(func,video_fpaths) - - print('Done!\n') + + print('Done!\n') \ No newline at end of file diff --git a/src/lib/utils.py b/src/lib/utils.py index b13be4a..b84e52e 100644 --- a/src/lib/utils.py +++ b/src/lib/utils.py @@ -146,7 +146,7 @@ def save_scene(out_fpath, k_arr, d_arr, r_arr, t_arr, cam_res): def save_optimised_cheetah(positions, out_fpath, extra_data=None, for_matlab=True, save_as_csv=True): file_data = dict(positions=positions) - + if extra_data is not None: assert type(extra_data) is dict file_data.update(extra_data) @@ -154,19 +154,19 @@ def save_optimised_cheetah(positions, out_fpath, extra_data=None, for_matlab=Tru with open(out_fpath, 'wb') as f: pickle.dump(file_data, f) print('Saved', out_fpath) - + if for_matlab: out_fpath = os.path.splitext(out_fpath)[0] + '.mat' savemat(out_fpath, file_data) print('Saved', out_fpath) - + if save_as_csv: # to-do?? # should use a similar method as save_3d_cheetah_as 3d, along the lines of: - + # xyz_labels = ['x', 'y', 'z'] # pdindex = pd.MultiIndex.from_product([bodyparts, xyz_labels], names=["bodyparts", "coords"]) - + # for i in range(len(video_fpaths)): # cam_name = os.path.splitext(os.path.basename(video_fpaths[i]))[0] # fpath = os.path.join(out_dir, cam_name + '_' + out_fname + '.h5') @@ -175,7 +175,7 @@ def save_optimised_cheetah(positions, out_fpath, extra_data=None, for_matlab=Tru # df.to_csv(os.path.splitext(fpath)[0] + ".csv") pass - + def save_3d_cheetah_as_2d(positions_3d, out_dir, scene_fpath, bodyparts, project_func, start_frame, save_as_csv=True, out_fname=None): assert os.path.dirname(os.path.dirname(scene_fpath)) in out_dir, 'scene_fpath does not belong to the same parent folder as out_dir' @@ -270,7 +270,7 @@ def get_pairwise_3d_points_from_df(points_2d_df, k_arr, d_arr, r_arr, t_arr, tri df_pairs = pd.concat([df_pairs, intersection_df], ignore_index=True, join='outer', sort=False) else: print(f"No pairwise points between camera {cam_a} and {cam_b}") - + print() points_3d_df = df_pairs[['frame', 'marker', 'x','y','z']].groupby(['frame','marker']).mean().reset_index() - return points_3d_df + return points_3d_df \ No newline at end of file diff --git a/src/lib/vid.py b/src/lib/vid.py index 753fe1d..fc898cb 100644 --- a/src/lib/vid.py +++ b/src/lib/vid.py @@ -71,7 +71,7 @@ def counter(self): def frame_count(self): return self.nframes - + def codec(self): return self.CODEC