Add support for passing a dict (cam_name -> files) into fit_eks_multicam

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',

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(

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'])