Add the option for multiple views loading - take2

movement/io/load_poses.py

@@ -351,6 +351,41 @@ def from_dlc_file(
     )
 
 
+def from_multiview_files(
+    file_path_dict: dict[str, Path | str],
+    source_software: Literal["DeepLabCut", "SLEAP", "LightningPose"],
+    fps: float | None = None,
+) -> xr.Dataset:
+    """Load and merge pose tracking data from multiple views (cameras).
+
+    Parameters
+    ----------
+    file_path_dict : dict[str, Union[Path, str]]
+        A dict whose keys are the view names and values are the paths to load.
+    source_software : {'LightningPose', 'SLEAP', 'DeepLabCut'}
+        The source software of the file.
+    fps : float, optional
+        The number of frames per second in the video. If None (default),
+        the `time` coordinates will be in frame numbers.
+
+    Returns
+    -------
+    xarray.Dataset
+        Dataset containing the pose tracks, confidence scores, and metadata,
+        with an additional views dimension.
+
+    """
+    views_list = list(file_path_dict.keys())
+    new_coord_views = xr.DataArray(views_list, dims="view")
+
+    dataset_list = [
+        from_file(f, source_software=source_software, fps=fps)
+        for f in file_path_dict.values()
+    ]
+
+    return xr.concat(dataset_list, dim=new_coord_views)
+
+
 def _ds_from_lp_or_dlc_file(
     file_path: Path | str,
     source_software: Literal["LightningPose", "DeepLabCut"],

tests/test_unit/test_load_poses.py

@@ -300,3 +300,20 @@ def test_from_numpy_valid(
             source_software=source_software,
         )
         self.assert_dataset(ds, expected_source_software=source_software)
+
+    def test_from_multiview_files(self):
+        """Test that the from_file() function delegates to the correct
+        loader function according to the source_software.
+        """
+        view_names = ["view_0", "view_1"]
+        file_path_dict = {
+            view: DATA_PATHS.get("DLC_single-wasp.predictions.h5")
+            for view in view_names
+        }
+        multi_view_ds = load_poses.from_multiview_files(
+            file_path_dict, source_software="DeepLabCut"
+        )
+
+        assert isinstance(multi_view_ds, xr.Dataset)
+        assert "view" in multi_view_ds.dims
+        assert multi_view_ds.view.values.tolist() == view_names