From d817d939d1198ca504872c9dcc2b019343d4fd10 Mon Sep 17 00:00:00 2001 From: Tuan Nguyen Date: Thu, 4 Jun 2026 09:26:47 -0400 Subject: [PATCH 1/3] Save clip_start_frame attribute to output pose HDF5 file Co-Authored-By: Claude Sonnet 4.6 --- src/mouse_tracking/utils/writers.py | 1 + tests/utils/writers/test_write_pose_clip.py | 23 +++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/src/mouse_tracking/utils/writers.py b/src/mouse_tracking/utils/writers.py index 5f3102c3..3b80f941 100644 --- a/src/mouse_tracking/utils/writers.py +++ b/src/mouse_tracking/utils/writers.py @@ -594,6 +594,7 @@ def write_pose_clip( for key, attrs in all_attrs.items(): for cur_attr, data in attrs.items(): out_f[key].attrs.create(cur_attr, data) + out_f["poseest"].attrs.create("clip_start_frame", clip_idxs[0]) def downgrade_pose_file(pose_h5_path, disable_id: bool = False): diff --git a/tests/utils/writers/test_write_pose_clip.py b/tests/utils/writers/test_write_pose_clip.py index 74868e44..be40b640 100644 --- a/tests/utils/writers/test_write_pose_clip.py +++ b/tests/utils/writers/test_write_pose_clip.py @@ -865,3 +865,26 @@ def test_comprehensive_pose_file_structure(): for file_path in [in_pose_file, out_pose_file]: if os.path.exists(file_path): os.unlink(file_path) + + +def test_writes_clip_start_frame_attribute(): + """Test that clip_start_frame attribute is written to the output pose file.""" + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp_in: + in_pose_file = tmp_in.name + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp_out: + out_pose_file = tmp_out.name + + try: + with h5py.File(in_pose_file, "w") as f: + poseest = f.create_group("poseest") + poseest.create_dataset("points", data=np.random.rand(20, 1, 12, 2).astype(np.float32)) + poseest.attrs["version"] = [3, 0] + + write_pose_clip(in_pose_file, out_pose_file, range(5, 15)) + + with h5py.File(out_pose_file, "r") as f: + assert f["poseest"].attrs["clip_start_frame"] == 5 + finally: + for path in [in_pose_file, out_pose_file]: + if os.path.exists(path): + os.unlink(path) From 1e05bac09c9dc74f825e295cef71d56463e1c2d0 Mon Sep 17 00:00:00 2001 From: Tuan Nguyen Date: Thu, 4 Jun 2026 10:21:03 -0400 Subject: [PATCH 2/3] Apply ruff formatting to test file Co-Authored-By: Claude Sonnet 4.6 --- tests/utils/writers/test_write_pose_clip.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/utils/writers/test_write_pose_clip.py b/tests/utils/writers/test_write_pose_clip.py index be40b640..9c8ff28d 100644 --- a/tests/utils/writers/test_write_pose_clip.py +++ b/tests/utils/writers/test_write_pose_clip.py @@ -877,7 +877,9 @@ def test_writes_clip_start_frame_attribute(): try: with h5py.File(in_pose_file, "w") as f: poseest = f.create_group("poseest") - poseest.create_dataset("points", data=np.random.rand(20, 1, 12, 2).astype(np.float32)) + poseest.create_dataset( + "points", data=np.random.rand(20, 1, 12, 2).astype(np.float32) + ) poseest.attrs["version"] = [3, 0] write_pose_clip(in_pose_file, out_pose_file, range(5, 15)) From fbc4eea3effe44dce08143a7618fb5c07cf8634d Mon Sep 17 00:00:00 2001 From: Tuan Nguyen Date: Thu, 4 Jun 2026 10:30:32 -0400 Subject: [PATCH 3/3] Guard clip_start_frame write when clip indices are empty Co-Authored-By: Claude Sonnet 4.6 --- src/mouse_tracking/utils/writers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mouse_tracking/utils/writers.py b/src/mouse_tracking/utils/writers.py index 3b80f941..e48dd6a5 100644 --- a/src/mouse_tracking/utils/writers.py +++ b/src/mouse_tracking/utils/writers.py @@ -594,7 +594,8 @@ def write_pose_clip( for key, attrs in all_attrs.items(): for cur_attr, data in attrs.items(): out_f[key].attrs.create(cur_attr, data) - out_f["poseest"].attrs.create("clip_start_frame", clip_idxs[0]) + if len(adjusted_clip_idxs) > 0: + out_f["poseest"].attrs.create("clip_start_frame", clip_idxs[0]) def downgrade_pose_file(pose_h5_path, disable_id: bool = False):