Support saving DLC 3D poses

Author: lochhh

movement/io/save_poses.py

@@ -37,14 +37,18 @@ def _ds_to_dlc_style_df(
     pandas.DataFrame
 
     """
-    # Concatenate the pose tracks and confidence scores into one array
-    tracks_with_scores = np.concatenate(
-        (
-            ds.position.data,
-            ds.confidence.data[:, np.newaxis, ...],
-        ),
-        axis=1,
-    )
+    is_3d = "z" in columns.get_level_values("coords")
+    if is_3d:
+        tracks_with_scores = ds.position.data
+    else:
+        # Concatenate the pose tracks and confidence scores into one array
+        tracks_with_scores = np.concatenate(
+            (
+                ds.position.data,
+                ds.confidence.data[:, np.newaxis, ...],
+            ),
+            axis=1,
+        )
     # Reverse the order of the dimensions except for the time dimension
     transpose_order = [0] + list(range(tracks_with_scores.ndim - 1, 0, -1))
     tracks_with_scores = tracks_with_scores.transpose(transpose_order)
@@ -121,7 +125,12 @@ def to_dlc_style_df(
     _validate_dataset(ds, ValidPosesDataset)
     scorer = ["movement"]
     bodyparts = ds.coords["keypoints"].data.tolist()
-    coords = ds.coords["space"].data.tolist() + ["likelihood"]
+    base_coords = ds.coords["space"].data.tolist()
+    coords = (
+        base_coords
+        if "z" in ds.coords["space"]
+        else base_coords + ["likelihood"]
+    )
     individuals = ds.coords["individuals"].data.tolist()
 
     if split_individuals:

tests/test_integration/test_io.py

@@ -13,13 +13,17 @@ def dlc_output_file(request, tmp_path):
     return tmp_path / request.param
 
 
-def test_load_and_save_to_dlc_style_df(valid_dlc_poses_df):
+@pytest.mark.parametrize(
+    "dlc_poses_df", ["valid_dlc_poses_df", "valid_dlc_3d_poses_df"]
+)
+def test_load_and_save_to_dlc_style_df(dlc_poses_df, request):
     """Test that loading pose tracks from a DLC-style DataFrame and
     converting back to a DataFrame returns the same data values.
     """
-    ds = load_poses.from_dlc_style_df(valid_dlc_poses_df)
+    dlc_poses_df = request.getfixturevalue(dlc_poses_df)
+    ds = load_poses.from_dlc_style_df(dlc_poses_df)
     df = save_poses.to_dlc_style_df(ds, split_individuals=False)
-    np.testing.assert_allclose(df.values, valid_dlc_poses_df.values)
+    np.testing.assert_allclose(df.values, dlc_poses_df.values)
 
 
 def test_save_and_load_dlc_file(dlc_output_file, valid_poses_dataset):