Skip to content

Commit ea8d2be

Browse files
talmoclaude
andauthored
Add metadata support to SuggestionFrame for preserving group information (#251)
* Add metadata support to SuggestionFrame for preserving group information This commit adds a metadata dictionary to SuggestionFrame (similar to Video.backend_metadata) to store arbitrary metadata that isn't explicitly represented in the data model. The primary use case is preserving the "group" key when reading/writing SLP files, which was previously being discarded. Changes: - Add metadata attribute to SuggestionFrame class with factory default - Update read_suggestions to extract and pass group metadata - Update write_suggestions to write group metadata if available - Add comprehensive tests for metadata round-trip and backward compatibility This maintains backward compatibility while enabling metadata preservation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Add missing coverage test from previous PR * Fix linting issues in test_video_reading.py Remove unused variable and apply formatting changes. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Simplify group metadata handling with .get() accessor Use .get() with default value instead of conditional check since group key should always be present. This simplifies the code and improves coverage. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude <noreply@anthropic.com>
1 parent f8b681d commit ea8d2be

4 files changed

Lines changed: 98 additions & 2 deletions

File tree

sleap_io/io/slp.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -845,10 +845,14 @@ def read_suggestions(labels_path: str, videos: list[Video]) -> list[SuggestionFr
845845
suggestions = [json.loads(x) for x in suggestions]
846846
suggestions_objects = []
847847
for suggestion in suggestions:
848+
# Extract metadata (e.g., "group")
849+
metadata = {"group": suggestion.get("group", 0)}
850+
848851
suggestions_objects.append(
849852
SuggestionFrame(
850853
video=videos[int(suggestion["video"])],
851854
frame_idx=suggestion["frame_idx"],
855+
metadata=metadata,
852856
)
853857
)
854858
return suggestions_objects
@@ -864,13 +868,15 @@ def write_suggestions(
864868
suggestions: A list of `SuggestionFrame` objects to store the metadata for.
865869
videos: A list of `Video` objects.
866870
"""
867-
GROUP = 0 # TODO: Handle storing extraneous metadata.
868871
suggestions_json = []
869872
for suggestion in suggestions:
873+
# Get group from metadata if available, otherwise use default
874+
group = suggestion.metadata.get("group", 0) if suggestion.metadata else 0
875+
870876
suggestion_dict = {
871877
"video": str(videos.index(suggestion.video)),
872878
"frame_idx": suggestion.frame_idx,
873-
"group": GROUP,
879+
"group": group,
874880
}
875881
suggestion_json = np.bytes_(json.dumps(suggestion_dict, separators=(",", ":")))
876882
suggestions_json.append(suggestion_json)

sleap_io/model/suggestions.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,11 @@ class SuggestionFrame:
1414
Attributes:
1515
video: The video associated with the frame.
1616
frame_idx: The index of the frame in the video.
17+
metadata: Dictionary containing additional metadata that is not explicitly
18+
represented in the data model. This is used to store arbitrary metadata
19+
such as the "group" key when reading/writing SLP files.
1720
"""
1821

1922
video: Video
2023
frame_idx: int
24+
metadata: dict[str, any] = attrs.field(factory=dict)

tests/io/test_slp.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -866,6 +866,40 @@ def test_suggestions(tmpdir):
866866
assert len(loaded_suggestions) == 0
867867

868868

869+
def test_suggestions_metadata(tmpdir):
870+
"""Test that suggestion metadata (e.g., group) is preserved during read/write."""
871+
labels = Labels()
872+
labels.videos.append(Video.from_filename("fake.mp4"))
873+
874+
# Create suggestions with different group values in metadata
875+
labels.suggestions.append(
876+
SuggestionFrame(video=labels.video, frame_idx=0, metadata={"group": 0})
877+
)
878+
labels.suggestions.append(
879+
SuggestionFrame(video=labels.video, frame_idx=1, metadata={"group": 1})
880+
)
881+
labels.suggestions.append(
882+
SuggestionFrame(video=labels.video, frame_idx=2, metadata={"group": 2})
883+
)
884+
885+
# Write and read suggestions
886+
write_suggestions(tmpdir / "test.slp", labels.suggestions, labels.videos)
887+
loaded_suggestions = read_suggestions(tmpdir / "test.slp", labels.videos)
888+
889+
# Verify metadata is preserved
890+
assert len(loaded_suggestions) == 3
891+
assert loaded_suggestions[0].metadata["group"] == 0
892+
assert loaded_suggestions[1].metadata["group"] == 1
893+
assert loaded_suggestions[2].metadata["group"] == 2
894+
895+
# Test backward compatibility: suggestions without metadata default to group 0
896+
suggestion_no_metadata = SuggestionFrame(video=labels.video, frame_idx=3)
897+
write_suggestions(tmpdir / "test2.slp", [suggestion_no_metadata], labels.videos)
898+
loaded_suggestions = read_suggestions(tmpdir / "test2.slp", labels.videos)
899+
assert len(loaded_suggestions) == 1
900+
assert loaded_suggestions[0].metadata["group"] == 0
901+
902+
869903
def test_pkg_roundtrip(tmpdir, slp_minimal_pkg):
870904
labels = read_labels(slp_minimal_pkg)
871905
assert type(labels.video.backend) is HDF5Video

tests/io/test_video_reading.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -752,6 +752,31 @@ def test_plugin_name_normalization():
752752
normalize_plugin_name("invalid_plugin")
753753

754754

755+
def test_image_plugin_name_normalization():
756+
"""Test image plugin name normalization with various aliases."""
757+
from sleap_io.io.video_reading import normalize_image_plugin_name
758+
759+
# Test opencv aliases
760+
assert normalize_image_plugin_name("opencv") == "opencv"
761+
assert normalize_image_plugin_name("OpenCV") == "opencv"
762+
assert normalize_image_plugin_name("cv") == "opencv"
763+
assert normalize_image_plugin_name("cv2") == "opencv"
764+
assert normalize_image_plugin_name("CV2") == "opencv"
765+
assert normalize_image_plugin_name("ocv") == "opencv"
766+
767+
# Test imageio aliases
768+
assert normalize_image_plugin_name("imageio") == "imageio"
769+
assert normalize_image_plugin_name("iio") == "imageio"
770+
771+
# Test invalid plugin
772+
with pytest.raises(ValueError, match="Unknown image plugin"):
773+
normalize_image_plugin_name("invalid_plugin")
774+
775+
# Test invalid plugin that's valid for video but not images
776+
with pytest.raises(ValueError, match="Unknown image plugin"):
777+
normalize_image_plugin_name("pyav")
778+
779+
755780
def test_global_default_plugin():
756781
"""Test global default plugin functionality."""
757782
import sleap_io as sio
@@ -915,3 +940,30 @@ def test_image_video_plugin_with_grayscale(centered_pair_frame_paths):
915940
np.testing.assert_array_equal(frame_opencv, frame_imageio)
916941
assert frame_opencv.ndim == 3 # Always 3D (H, W, C)
917942
assert frame_opencv.shape[-1] in (1, 3) # Grayscale or RGB
943+
944+
945+
def test_image_video_default_plugin_without_opencv(
946+
centered_pair_frame_paths, monkeypatch
947+
):
948+
"""Test ImageVideo defaults to imageio when opencv not available."""
949+
import sys
950+
951+
# Mock sys.modules to simulate opencv not being available
952+
if "cv2" in sys.modules:
953+
monkeypatch.delitem(sys.modules, "cv2")
954+
955+
# Clear any global default
956+
import sleap_io as sio
957+
958+
original_default = sio.get_default_image_plugin()
959+
try:
960+
sio.set_default_image_plugin(None)
961+
962+
# Create ImageVideo without specifying plugin
963+
backend = ImageVideo(centered_pair_frame_paths)
964+
965+
# Should default to imageio since opencv is not available
966+
assert backend.plugin == "imageio"
967+
finally:
968+
# Restore
969+
sio.set_default_image_plugin(original_default)

0 commit comments

Comments
 (0)