Skip to content

Commit

Permalink
Update labels videos list on replace
Browse files Browse the repository at this point in the history
  • Loading branch information
talmo committed Oct 4, 2024
1 parent ddf5277 commit a58fbf0
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 3 deletions.
10 changes: 10 additions & 0 deletions sleap_io/model/labels.py
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,13 @@ def replace_videos(
video_map: Alternative input of dictionary where keys are the old videos and
values are the new videos.
"""
if (
old_videos is None
and new_videos is not None
and len(new_videos) == len(self.videos)
):
old_videos = self.videos

if video_map is None:
video_map = {o: n for o, n in zip(old_videos, new_videos)}

Expand All @@ -486,6 +493,9 @@ def replace_videos(
if sf.video in video_map:
sf.video = video_map[sf.video]

# Update the list of videos.
self.videos = [video_map.get(video, video) for video in self.videos]

def replace_filenames(
self,
new_filenames: list[str | Path] | None = None,
Expand Down
6 changes: 3 additions & 3 deletions tests/model/test_labels.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,16 +356,16 @@ def test_labels_remove_predictions(slp_real_data):
def test_replace_videos(slp_real_data):
labels = load_slp(slp_real_data)
assert labels.video.filename == "tests/data/videos/centered_pair_low_quality.mp4"
labels.replace_videos(
old_videos=[labels.video], new_videos=[Video.from_filename("fake.mp4")]
)
labels.replace_videos(new_videos=[Video.from_filename("fake.mp4")])

for lf in labels:
assert lf.video.filename == "fake.mp4"

for sf in labels.suggestions:
assert sf.video.filename == "fake.mp4"

assert labels.video.filename == "fake.mp4"


def test_replace_filenames():
labels = Labels(videos=[Video.from_filename("a.mp4"), Video.from_filename("b.mp4")])
Expand Down

0 comments on commit a58fbf0

Please sign in to comment.