diff --git a/sleap_io/io/slp.py b/sleap_io/io/slp.py index 4d31a577..1f172fe4 100644 --- a/sleap_io/io/slp.py +++ b/sleap_io/io/slp.py @@ -20,7 +20,7 @@ LabeledFrame, Labels, ) -from sleap_io.io.video import MediaVideo, HDF5Video +from sleap_io.io.video import ImageVideo, MediaVideo, HDF5Video from sleap_io.io.utils import ( read_hdf5_attrs, read_hdf5_dataset, @@ -72,16 +72,23 @@ def read_videos(labels_path: str) -> list[Video]: # complex path finding strategies. pass + video_path = video_path.as_posix() + + if "filenames" in backend: + # This is an ImageVideo. + # TODO: Path resolution. + video_path = backend["filenames"] + try: backend = VideoBackend.from_filename( - video_path.as_posix(), + video_path, dataset=backend.get("dataset", None), grayscale=backend.get("grayscale", None), input_format=backend.get("input_format", None), ) except ValueError: backend = None - video_objects.append(Video(filename=video_path.as_posix(), backend=backend)) + video_objects.append(Video(filename=video_path, backend=backend)) return video_objects @@ -119,6 +126,24 @@ def write_videos(labels_path: str, videos: list[Video]): # TODO: Handle saving embedded images or restoring source video. # Ref: https://github.com/talmolab/sleap/blob/fb61b6ce7a9ac9613d99303111f3daafaffc299b/sleap/io/format/hdf5.py#L246-L273 + elif type(video.backend) == ImageVideo: + shape = video.shape + if shape is None: + height, width, channels = 0, 0, 1 + else: + height, width, channels = shape[1:] + + video_json = { + "backend": { + "filename": video.filename[0], + "filenames": video.filename, + "height_": height, + "width_": width, + "channels_": channels, + "grayscale": video.backend.grayscale, + } + } + else: raise NotImplementedError( f"Cannot serialize video backend for video: {video}" diff --git a/tests/data/slp/imgvideo.slp b/tests/data/slp/imgvideo.slp new file mode 100644 index 00000000..0369499c Binary files /dev/null and b/tests/data/slp/imgvideo.slp differ diff --git a/tests/fixtures/slp.py b/tests/fixtures/slp.py index 1b82a32e..b9ec1f4c 100644 --- a/tests/fixtures/slp.py +++ b/tests/fixtures/slp.py @@ -70,3 +70,9 @@ def slp_real_data(): "/Users/talmo/sleap-io/tests/data/videos/centered_pair_low_quality.mp4" """ return "tests/data/slp/labels.v002.rel_paths.slp" + + +@pytest.fixture +def slp_imgvideo(): + """SLP project with a single image video.""" + return "tests/data/slp/imgvideo.slp" diff --git a/tests/io/test_slp.py b/tests/io/test_slp.py index 399674c3..e2ab9470 100644 --- a/tests/io/test_slp.py +++ b/tests/io/test_slp.py @@ -33,6 +33,8 @@ from sleap_io.io.utils import read_hdf5_dataset import numpy as np +from sleap_io.io.video import ImageVideo + def test_read_labels(slp_typical, slp_simple_skel, slp_minimal): """Test `read_labels` can read different types of .slp files.""" @@ -215,3 +217,23 @@ def test_load_multi_skeleton(tmpdir): assert loaded_skels[1].edge_inds == [(0, 1)] assert loaded_skels[0].flipped_node_inds == [1, 0] assert loaded_skels[1].flipped_node_inds == [1, 0] + + +def test_slp_imgvideo(tmpdir, slp_imgvideo): + labels = read_labels(slp_imgvideo) + assert type(labels.video.backend) == ImageVideo + assert labels.video.shape == (3, 384, 384, 1) + + write_labels(tmpdir / "test.slp", labels) + labels = read_labels(tmpdir / "test.slp") + assert type(labels.video.backend) == ImageVideo + assert labels.video.shape == (3, 384, 384, 1) + + videos = [Video.from_filename(["fake1.jpg", "fake2.jpg"])] + assert videos[0].shape is None + assert len(videos[0].filename) == 2 + write_videos(tmpdir / "test2.slp", videos) + videos = read_videos(tmpdir / "test2.slp") + assert type(videos[0].backend) == ImageVideo + assert len(videos[0].filename) == 2 + assert videos[0].shape is None