Adding support for more .nwb training data to SLEAP

README.md

@@ -41,6 +41,11 @@ labels = sio.load_file("predictions.slp")
 sio.save_file(labels, "predictions.nwb")
 # Or:
 # labels.save("predictions.nwb")
+
+# Save to an NWB file and convert SLEAP training data to NWB training data.
+# Note: This will save the first 3 frames of the video as individual images in a
+# subfolder next to the NWB file for portability.
+sio.save_file(labels, "predictions.nwb", as_training=True, frame_inds=[0, 1, 2])
 ```
 
 ### Convert labels to raw arrays

docs/index.md

@@ -45,6 +45,11 @@ labels = sio.load_file("predictions.slp")
 
 # Save to NWB file.
 labels.save("predictions.nwb")
+
+# Save to an NWB file and convert SLEAP training data to NWB training data.
+# Note: This will save the first 3 frames of the video as individual images in a
+# subfolder next to the NWB file for portability.
+sio.save_file(labels, "predictions.nwb", as_training=True, frame_inds=[0, 1, 2])
 ```
 
 **See also:** [`Labels.save`](model.md#sleap_io.Labels.save) and [Formats](formats.md)

environment.yml

@@ -24,7 +24,7 @@ dependencies:
   - pydocstyle
   - toml
   - twine
-  - build
+  - python-build
   - pip
   - pip:
     - "--editable=.[dev]"

pyproject.toml

@@ -24,7 +24,9 @@ dependencies = [
     "attrs",
     "h5py>=3.8.0",
     "pynwb",
-    "ndx-pose",
+    # TODO: Replace with ndx-pose>=0.2.0 once released.
+    # https://github.com/rly/ndx-pose/issues/29
+    "ndx-pose @ git+https://github.com/rly/ndx-pose@a847ad4be75e60ef9e413b8cbfc99c616fc9fd05",
     "pandas",
     "simplejson",
     "imageio",

sleap_io/io/main.py

@@ -59,21 +59,57 @@ def load_nwb(filename: str) -> Labels:
     return nwb.read_nwb(filename)
 
 
-def save_nwb(labels: Labels, filename: str, append: bool = True):
+def save_nwb(
+    labels: Labels,
+    filename: str,
+    as_training: bool = False,
+    append: bool = True,
+    frame_inds: Optional[list[int]] = None,
+    frame_path: Optional[str] = None,
+    image_format: str = "png",
+):
     """Save a SLEAP dataset to NWB format.
 
     Args:
         labels: A SLEAP `Labels` object (see `load_slp`).
         filename: Path to NWB file to save to. Must end in `.nwb`.
+        as_training: If `True`, save the dataset as a training dataset. This will use
+            data structures from `ndx-pose` that do not assume the data is a continuous
+            timeseries. This is useful for saving training data which includes frames
+            from random videos. If `False` (the default), save the dataset as a
+            timeseries, assuming the data is continuous. This is useful for saving
+            tracked predictions that is used for downstream analysis.
         append: If `True` (the default), append to existing NWB file. File will be
             created if it does not exist.
-
-    See also: nwb.write_nwb, nwb.append_nwb
+        frame_inds: Optional list of labeled frame indices within the Labels to save
+            when saving in training data format. If `None`, all labeled frames in the
+            labels will be saved. No effect if `as_training` is `False`.
+        frame_path: The path to a folder to save the extracted frame images to when
+            saving in training data format. If `None`, the path is the NWB filename
+            without the extension. No effect if `as_training` is `False`.
+        image_format: The image format to use when saving extracted frame images.
+            Defaults to "png". No effect if `as_training` is `False`.
+
+    See also: nwb.write_nwb, nwb.append_nwb, nwb.append_nwb_training
     """
     if append and Path(filename).exists():
-        nwb.append_nwb(labels, filename)
+        nwb.append_nwb(
+            labels,
+            filename,
+            as_training=as_training,
+            frame_inds=frame_inds,
+            frame_path=frame_path,
+            image_format=image_format,
+        )
     else:
-        nwb.write_nwb(labels, filename)
+        nwb.write_nwb(
+            labels,
+            filename,
+            as_training=as_training,
+            frame_inds=frame_inds,
+            frame_path=frame_path,
+            image_format=image_format,
+        )
 
 
 def load_labelstudio(
@@ -190,6 +226,8 @@ def load_file(
         return load_jabs(filename, **kwargs)
     elif format == "video":
         return load_video(filename, **kwargs)
+    else:
+        raise ValueError(f"Unknown format '{format}' for filename: '{filename}'.")
 
 
 def save_file(
@@ -219,8 +257,10 @@ def save_file(
 
     if format == "slp":
         save_slp(labels, filename, **kwargs)
-    elif format == "nwb":
-        save_nwb(labels, filename, **kwargs)
+    elif format in ("nwb", "nwb_predictions"):
+        save_nwb(labels, filename, False)
+    elif format == "nwb_training":
+        save_nwb(labels, filename, True, frame_inds=kwargs.get("frame_inds", None))
     elif format == "labelstudio":
         save_labelstudio(labels, filename, **kwargs)
     elif format == "jabs":