Skip to content

Commit

Permalink
n
Browse files Browse the repository at this point in the history
  • Loading branch information
keyaloding committed Jul 16, 2024
1 parent 23a5a83 commit cf4bcf7
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions sleap_io/io/nwb.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import datetime
import uuid
import re
from imageio.v3 import imwrite

import pandas as pd
import numpy as np
Expand Down Expand Up @@ -216,6 +217,25 @@ def videos_to_source_videos(videos: List[Video]) -> SourceVideos: # type: ignor
return SourceVideos(image_series=source_videos)


def sleap_pkg_to_nwb(filename: str, labels: Labels, **kwargs):
"""Write a SLEAP package to an NWB file.
Args:
filename: The path to the SLEAP package.
labels: The SLEAP Labels object.
"""
assert filename.endswith(".pkg.slp")

path = filename.split(".slp")[0]
save_path = Path(path + ".nwb_images")
img_paths = []
for i, labeled_frame in enumerate(labels.labeled_frames):
img_path = save_path / f"frame_{i}.png"
imwrite(img_path, labeled_frame.image)
img_paths.append(img_path)



def get_timestamps(series: PoseEstimationSeries) -> np.ndarray:
"""Return a vector of timestamps for a `PoseEstimationSeries`."""
if series.timestamps is not None:
Expand Down

0 comments on commit cf4bcf7

Please sign in to comment.