diff --git a/eks/multicam_smoother.py b/eks/multicam_smoother.py index 1cdb124..e43503d 100644 --- a/eks/multicam_smoother.py +++ b/eks/multicam_smoother.py @@ -154,7 +154,7 @@ def fit_eks_mirrored_multicam( def fit_eks_multicam( - input_source: str | list, + input_source: str | list | dict, save_dir: str, bodypart_list: list | None = None, smooth_param: float | list | None = None, @@ -172,7 +172,9 @@ def fit_eks_multicam( Fit the Ensemble Kalman Smoother for un-mirrored multi-camera data. Args: - input_source: Directory path or list of CSV file paths with columns for all cameras. + input_source: Directory path, a list of CSV file paths with columns for all cameras, + or a dictionary mapping each camera name to a list of CSV file paths. + (camera names must match either camera_names or the calibration .toml). save_dir: Directory to save output DataFrame. bodypart_list: List of body parts. smooth_param: Value in (0, Inf); smaller values lead to more smoothing. @@ -209,6 +211,14 @@ def fit_eks_multicam( ... camera_names=['cam0', 'cam1', 'cam2'], ... ) + >>> from eks import fit_eks_multicam + >>> camera_dfs, smooth_params, input_dfs, bodyparts, df_3d = fit_eks_multicam( + ... input_source={'cam0': ['/path/to/seed0/cam0.csv', '/path/to/seed1/cam0.csv'], + 'cam1': ['/path/to/seed0/cam1.csv', '/path/to/seed1/cam1.csv']}, + ... save_dir='/path/to/output/', + ... camera_names=['cam0', 'cam1'] + ... ) + >>> # with a calibration file for nonlinear 3D projection >>> camera_dfs, smooth_params, input_dfs, bodyparts, df_3d = fit_eks_multicam( ... input_source='/path/to/csv/dir', diff --git a/eks/utils.py b/eks/utils.py index b7211f7..2de5902 100644 --- a/eks/utils.py +++ b/eks/utils.py @@ -135,12 +135,16 @@ def get_keypoint_names(df: pd.DataFrame) -> list: return kps.tolist() -def format_data(input_source: str | list, camera_names: list | None = None) -> tuple[list, list]: +def format_data( + input_source: str | list | dict, + camera_names: list | None = None, +) -> tuple[list, list]: """ Load and format input files from a directory or a list of file paths. Args: - input_source (str or list): Directory path or list of file paths. + input_source (str or list or dict): Directory path, list of file paths, or dictionary + mapping camera_names to list of file paths per camera. camera_names (None or list): List of multiple camera/view names. None = single camera *** data with mirrored naming schemes (e.g. paw1LH_top), keep camera_names as None @@ -161,8 +165,16 @@ def format_data(input_source: str | list, camera_names: list | None = None) -> t elif isinstance(input_source, list): # If it's a list of file paths, use it directly file_paths = sorted(input_source) + elif isinstance(input_source, dict): + # If it's a dictionary, keep as is, and we pick the correct entry below + # when looping over camera_names + file_paths = input_source + else: - raise ValueError('input_source must be a directory path or a list of file paths') + raise ValueError( + 'input_source must be a directory path, ' + 'a list of file paths, or a map from camera names to list of file paths' + ) # Process each file based on the data type if camera_names is None: @@ -181,7 +193,8 @@ def format_data(input_source: str | list, camera_names: list | None = None) -> t input_dfs_list.append(markers_curr_fmt) else: for camera in camera_names: - matched = [fp for fp in file_paths if camera in os.path.basename(fp)] + files = file_paths if isinstance(file_paths, list) else file_paths.get(camera, []) + matched = [fp for fp in files if camera in os.path.basename(fp)] valid = [fp for fp in matched if fp.endswith('.csv') or fp.endswith('.slp')] if len(valid) == 0: raise FileNotFoundError( diff --git a/tests/test_utils.py b/tests/test_utils.py index 0c5caac..0c4d3cc 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -191,6 +191,24 @@ def test_format_data_camera_names_multiple_models(self, tmp_path): assert len(input_dfs_list[0]) == 2 assert len(input_dfs_list[1]) == 2 + def test_format_data_camera_mapping_input(self, tmp_path): + """Test loading camera-specific files from a mapping.""" + # Arrange + top0 = _make_dlc_csv(tmp_path, 'model0_top.csv', ['nose']) + top1 = _make_dlc_csv(tmp_path, 'model1_top.csv', ['nose']) + bot0 = _make_dlc_csv(tmp_path, 'model0_bot.csv', ['nose']) + bot1 = _make_dlc_csv(tmp_path, 'model1_bot.csv', ['nose']) + + # Act + input_dfs_list, keypoint_names = format_data( + {'top': [str(top1), str(top0)], 'bot': [str(bot1), str(bot0)]}, + camera_names=['top', 'bot'], + ) + + # Assert + assert [len(dfs) for dfs in input_dfs_list] == [2, 2] + assert keypoint_names == ['nose'] + def test_format_data_camera_not_found_raises(self, tmp_path): # Arrange - only 'top' files exist _make_dlc_csv(tmp_path, 'model0_top.csv', ['nose'])