Skip to content

Commit

Permalink
Implement ImageVideo support in SLP (#90)
Browse files Browse the repository at this point in the history
* Add fixture

* Read, write and tests

* coverage
  • Loading branch information
talmo authored May 1, 2024
1 parent 5323396 commit e553986
Show file tree
Hide file tree
Showing 4 changed files with 56 additions and 3 deletions.
31 changes: 28 additions & 3 deletions sleap_io/io/slp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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}"
Expand Down
Binary file added tests/data/slp/imgvideo.slp
Binary file not shown.
6 changes: 6 additions & 0 deletions tests/fixtures/slp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
22 changes: 22 additions & 0 deletions tests/io/test_slp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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

0 comments on commit e553986

Please sign in to comment.