From 27bd8a5f3f86b8de853414e2ed2b426f6a0ee272 Mon Sep 17 00:00:00 2001
From: ishan372or <ishan372or@gmail.com>
Date: Mon, 1 Jun 2026 21:00:05 +0530
Subject: [PATCH] Handle individual-wise confidence during DLC/LP export

---
 movement/io/save_poses.py                  | 23 ++++++++-
 tests/test_unit/test_io/test_save_poses.py | 54 ++++++++++++++++++++++
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/movement/io/save_poses.py b/movement/io/save_poses.py
index a9c0a5b5e..1db196c8c 100644
--- a/movement/io/save_poses.py
+++ b/movement/io/save_poses.py
@@ -39,13 +39,34 @@ def _ds_to_dlc_style_df(
     """
     # Keep position data as is, if data is 3D (i.e. contains 'z' coordinate)
     # Otherwise, concatenate position and confidence scores into one array
+    # DLC stores confidence scores per keypoint. If confidence is provided
+    # per individual, expand it across all keypoints before export.
+    confidence = ds.confidence.data
+    if "keypoint" not in ds.confidence.dims:
+        logger.warning(
+            "Dataset contains individual-wise confidence scores. "
+            "DeepLabCut only supports keypoint-wise confidence scores, "
+            "so confidence values will be expanded to all keypoints."
+        )
+        if confidence.ndim == 1:
+            confidence = np.repeat(
+                confidence[:, np.newaxis],
+                ds.sizes["keypoint"],
+                axis=1,
+            )
+        else:
+            confidence = np.repeat(
+                confidence[:, np.newaxis, :],
+                ds.sizes["keypoint"],
+                axis=1,
+            )
     tracks = (
         ds.position.data
         if "z" in columns.get_level_values("coords")
         else np.concatenate(
             (
                 ds.position.data,
-                ds.confidence.data[:, np.newaxis, ...],
+                confidence[:, np.newaxis, ...],
             ),
             axis=1,
         )
diff --git a/tests/test_unit/test_io/test_save_poses.py b/tests/test_unit/test_io/test_save_poses.py
index 5c465dd44..2358bbd6d 100644
--- a/tests/test_unit/test_io/test_save_poses.py
+++ b/tests/test_unit/test_io/test_save_poses.py
@@ -118,6 +118,44 @@ def test_to_dlc_style_df(ds, expected_exception):
             ]
 
 
+def test_to_dlc_style_df_with_individual_confidence():
+    """Test that datasets with individual-wise confidence scores can
+    be converted to a DLC-style DataFrame.
+    """
+    ds = load_poses.from_dlc_file(
+        DATA_PATHS.get("DLC_single-wasp.predictions.h5")
+    )
+
+    # Convert keypoint-wise confidence to individual-wise confidence
+    ds["confidence"] = ds.confidence.isel(keypoint=0)
+
+    df = save_poses.to_dlc_style_df(
+        ds,
+        split_individuals=False,
+    )
+
+    assert isinstance(df, pd.DataFrame)
+
+
+def test_to_dlc_file_with_individual_confidence(tmp_path):
+    """Test that datasets with individual-wise confidence scores can
+    be exported to DLC format.
+    """
+    ds = load_poses.from_dlc_file(
+        DATA_PATHS.get("DLC_single-wasp.predictions.h5")
+    )
+
+    ds["confidence"] = ds.confidence.isel(keypoint=0)
+
+    save_poses.to_dlc_file(
+        ds,
+        tmp_path / "test.h5",
+        split_individuals=False,
+    )
+
+    assert (tmp_path / "test.h5").is_file()
+
+
 def test_to_dlc_file_valid_dataset(
     output_file_params, valid_poses_dataset, request
 ):
@@ -275,6 +313,22 @@ def test_to_lp_file_invalid_dataset(
         )
 
 
+def test_to_lp_file_with_individual_confidence(tmp_path):
+    """Test that datasets with individual-wise confidence scores can
+    be exported to LightningPose format.
+    """
+    ds = load_poses.from_dlc_file(
+        DATA_PATHS.get("LP_mouse-face_AIND.predictions.csv")
+    )
+
+    ds["confidence"] = ds.confidence.isel(keypoint=0)
+
+    save_poses.to_lp_file(
+        ds,
+        tmp_path / "test.csv",
+    )
+
+
 def test_to_sleap_analysis_file_valid_dataset(
     output_file_params, valid_poses_dataset, request
 ):